From 4392b4fee156f3b975d2650de3bb98a095318011 Mon Sep 17 00:00:00 2001 From: pphua Date: Mon, 24 Aug 2026 19:29:00 +0800 Subject: [PATCH] feat(tasks): add runtime controls and inherited context --- agent/blackboard_inheritance_test.go | 274 ++++++++ agent/cancelcause.go | 90 +++ agent/capture.go | 58 +- agent/chat.go | 1 + agent/goals.go | 21 +- agent/mainagent.go | 24 +- agent/planner.go | 25 +- agent/runinfo.go | 47 ++ agent/terminalreason.go | 161 +++++ agent/terminalreason_test.go | 169 +++++ agent/tools.go | 477 +++++++++----- agent/tools_insert.go | 44 +- agent/tools_overview_test.go | 31 + agent/worker.go | 29 +- cmd/artex/main.go | 20 +- cmd/artex/main_test.go | 19 + db/db.go | 4 +- db/exploration.go | 366 +++++++++-- db/exploration_sources.go | 350 ++++++++++ db/exploration_sources_test.go | 431 +++++++++++++ db/exploration_test.go | 136 +++- db/exploration_tokens_test.go | 130 ++++ db/intent_admission.go | 32 + db/intent_control_test.go | 285 +++++++++ db/task_assets_context.go | 20 +- db/task_context.go | 419 ++++++++++++ db/task_context_lock_test.go | 527 ++++++++++++++++ db/task_context_unit_test.go | 27 + db/task_delete_concurrency_test.go | 210 ++++++ db/task_queue_test.go | 66 ++ db/task_source_limit_test.go | 47 ++ db/tasks.go | 326 ++++++++-- db/tasks_test.go | 507 ++++++++++++++- server/chatupload.go | 13 + server/engine.go | 660 +++++++++++++++---- server/engine_cancelcause_test.go | 217 +++++++ server/engine_timeout.go | 70 +- server/goals.go | 389 ++++++++++-- server/goals_test.go | 12 + server/inheritance_api_test.go | 100 +++ server/inheritance_dto_test.go | 117 ++++ server/manager.go | 913 +++++++++++++++++++++++++-- server/manager_delete_files_test.go | 280 ++++++++ server/manager_lifecycle_test.go | 109 ++++ server/orchestration.go | 23 +- server/scheduler.go | 39 +- server/task_control_routes_test.go | 46 ++ server/task_delete_barrier_test.go | 364 +++++++++++ traffic/traffic.go | 65 +- traffic/traffic_test.go | 164 +++++ 50 files changed, 8291 insertions(+), 663 deletions(-) create mode 100644 agent/blackboard_inheritance_test.go create mode 100644 agent/cancelcause.go create mode 100644 agent/runinfo.go create mode 100644 agent/terminalreason.go create mode 100644 agent/terminalreason_test.go create mode 100644 agent/tools_overview_test.go create mode 100644 cmd/artex/main_test.go create mode 100644 db/exploration_sources.go create mode 100644 db/exploration_sources_test.go create mode 100644 db/exploration_tokens_test.go create mode 100644 db/intent_admission.go create mode 100644 db/intent_control_test.go create mode 100644 db/task_context.go create mode 100644 db/task_context_lock_test.go create mode 100644 db/task_context_unit_test.go create mode 100644 db/task_delete_concurrency_test.go create mode 100644 db/task_queue_test.go create mode 100644 db/task_source_limit_test.go create mode 100644 server/engine_cancelcause_test.go create mode 100644 server/goals_test.go create mode 100644 server/inheritance_api_test.go create mode 100644 server/inheritance_dto_test.go create mode 100644 server/manager_delete_files_test.go create mode 100644 server/manager_lifecycle_test.go create mode 100644 server/task_control_routes_test.go create mode 100644 server/task_delete_barrier_test.go diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go new file mode 100644 index 0000000..343d422 --- /dev/null +++ b/agent/blackboard_inheritance_test.go @@ -0,0 +1,274 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/Autumn-27/artex/db" + actool "github.com/Autumn-27/norma/tool" +) + +func callReadJSON(t *testing.T, tool actool.CoreTool, input string) any { + t.Helper() + result, err := tool.Call(context.Background(), json.RawMessage(input), nil) + if err != nil { + t.Fatalf("tool call: %v", err) + } + var out any + if err := json.Unmarshal([]byte(result.Flatten()), &out); err != nil { + t.Fatalf("decode tool result: %v; raw=%s", err, result.Flatten()) + } + return out +} + +func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) { + d := testDB(t) + defer d.Close() + + companies := d.Companies() + companyID, _, err := companies.UpsertCompany(fmt.Sprintf("overview-scope-%d", time.Now().UnixNano()), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = companies.DeleteCompany(companyID) }) + + domain := fmt.Sprintf("overview-scope-%d.invalid", companyID) + ip := fmt.Sprintf("2001:db8:%x::42", companyID%0xffff) + cidr := fmt.Sprintf("2001:db8:%x:1::/64", companyID%0xffff) + icp := fmt.Sprintf("京 ICP 备 %d 号", companyID) + keyword := fmt.Sprintf("Scope Company %d", companyID) + inputs := []db.ScopeInput{ + {Kind: "domain", Value: domain}, + {Kind: "ip", Value: ip}, + {Kind: "cidr", Value: cidr}, + {Kind: "icp", Value: icp}, + {Kind: "keyword", Value: keyword}, + } + added, skipped, invalid, scopeErrors := companies.AddScopeInputs(companyID, inputs, "task context test") + if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 { + t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors) + } + + task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{ + CompanyIDs: []int64{companyID}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + + tools := NewToolSet(d.Exploration(task.ExplorationID), "planner") + tools.SetTaskID(task.ID) + tools.SetAssetStore(d.Assets(), companies) + overview := tools.graphOverviewData() + coverage, ok := overview["coverage"].(map[string]any) + if !ok { + t.Fatalf("coverage missing: %#v", overview["coverage"]) + } + scopeRows, ok := coverage["scope"].([]map[string]any) + if !ok || len(scopeRows) != 1 { + t.Fatalf("task scope missing: %#v", coverage["scope"]) + } + companyScope, ok := scopeRows[0]["company_scope"].([]map[string]any) + if !ok || len(companyScope) != len(inputs) { + t.Fatalf("company scope not expanded: %#v", scopeRows[0]) + } + kinds := make(map[string]string, len(companyScope)) + for _, rule := range companyScope { + kinds[fmt.Sprint(rule["kind"])] = fmt.Sprint(rule["value"]) + } + for _, input := range inputs { + if kinds[input.Kind] != input.Value { + t.Errorf("scope %s=%q want %q", input.Kind, kinds[input.Kind], input.Value) + } + } + keywords, ok := scopeRows[0]["company_keywords"].([]string) + if !ok || len(keywords) != 1 || keywords[0] != keyword { + t.Fatalf("company keywords missing: %#v", scopeRows[0]["company_keywords"]) + } +} + +func TestBlackboardToolsReadDirectSources(t *testing.T) { + d := testDB(t) + defer d.Close() + + grand, err := d.CreateTask("grand", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("source", "source goal", db.TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}}) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + }) + + grandStore := d.Exploration(grand.ExplorationID) + sourceStore := d.Exploration(source.ExplorationID) + currentStore := d.Exploration(current.ExplorationID) + grandFact, err := grandStore.AddNode(db.KindFact, map[string]any{"summary": "indirect-only"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFact, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "shared fact", "confidence": "verified"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "shared work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil { + t.Fatal(err) + } + sourceFinding, err := sourceStore.AddNode(db.KindFinding, map[string]any{ + "summary": "shared finding", "vulnclass": "idor", "severity": "high", + }, 9, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + if err := sourceStore.Link(sourceIntent, db.RelYields, sourceFinding); err != nil { + t.Fatal(err) + } + stepID, err := sourceStore.AppendActivity(db.Activity{ + NodeID: &sourceIntent, Worker: "source-worker", Kind: "result", Tool: "HTTP", + Summary: "shared trace marker", Detail: "shared trace full detail", + }) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 301; i++ { + if _, err := sourceStore.AddIntent(map[string]any{"summary": fmt.Sprintf("newer live source work %d", i)}, 1, nil, "planner"); err != nil { + t.Fatal(err) + } + } + assets := d.Assets() + host := fmt.Sprintf("overview-host-%d.invalid", source.ID) + assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) }) + if _, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "source host anchor"}, 0, "confirmed", "worker", []int64{assetID}); err != nil { + t.Fatal(err) + } + + tools := NewToolSet(currentStore, "worker") + tools.SetTaskID(current.ID) + tools.SetAssetStore(assets, assets.Companies()) + parentID, _ := json.Marshal(sourceFact) + derivedID, err := tools.addOneIntent(intentItem{ + Summary: "derive locally from shared fact", ParentIDs: []json.RawMessage{parentID}, + }) + if err != nil { + t.Fatalf("derive from inherited fact: %v", err) + } + if derived, err := currentStore.GetNode(derivedID); err != nil || derived == nil { + t.Fatalf("derived intent must be local: node=%+v err=%v", derived, err) + } + beforeFacts, _ := currentStore.ListByKind(db.KindFact, 100) + sourceIntentID, _ := json.Marshal(sourceIntent) + if _, err := tools.recordOneFact(factItem{Summary: "must not attach to inherited intent", IntentID: json.RawMessage(sourceIntentID)}, sourceIntent); err == nil { + t.Fatal("record_fact must reject inherited intent") + } + afterFacts, _ := currentStore.ListByKind(db.KindFact, 100) + if len(afterFacts) != len(beforeFacts) { + t.Fatalf("rejected inherited write persisted a fact: before=%d after=%d", len(beforeFacts), len(afterFacts)) + } + overview := tools.graphOverviewData() + related, ok := overview["related_tasks"].([]map[string]any) + if !ok || len(related) != 1 { + t.Fatalf("related_tasks: %#v", overview["related_tasks"]) + } + if related[0]["source_task_id"] != source.ID || related[0]["inherited"] != true { + t.Fatalf("related source provenance: %#v", related[0]) + } + recentFacts, ok := related[0]["recent_facts"].([]map[string]any) + foundSourceFact := false + for _, fact := range recentFacts { + if fact["id"] == sourceFact { + foundSourceFact = true + break + } + } + if !ok || !foundSourceFact { + t.Fatalf("related fact summary: %#v", related[0]["recent_facts"]) + } + if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding { + t.Fatalf("related finding summary: %#v", related[0]["recent_findings"]) + } + results, ok := related[0]["recent_intent_results"].([]map[string]any) + if !ok || len(results) == 0 || results[0]["id"] != sourceIntent { + t.Fatalf("terminal source intent was starved by newer live work: %#v", related[0]["recent_intent_results"]) + } + if results[0]["result_summary"] != "shared trace marker" { + t.Fatalf("terminal source intent result was not distilled: %#v", results[0]) + } + coverage, ok := overview["coverage"].(map[string]any) + if !ok { + t.Fatalf("coverage missing from overview: %#v", overview["coverage"]) + } + hosts, ok := coverage["hosts"].([]string) + if !ok || len(hosts) != 1 || hosts[0] != host || coverage["host_count"] != 1 { + t.Fatalf("inherited host context missing: %#v", coverage) + } + + facts := callReadJSON(t, tools.listFacts(), `{}`).([]any) + seenSource, seenGrand := false, false + for _, raw := range facts { + item := raw.(map[string]any) + id := int64(item["id"].(float64)) + if id == sourceFact { + seenSource = item["inherited"] == true && int64(item["source_task_id"].(float64)) == source.ID + } + seenGrand = seenGrand || id == grandFact + } + if !seenSource || seenGrand { + t.Fatalf("list_facts direct-only: source=%v grand=%v payload=%#v", seenSource, seenGrand, facts) + } + + findings := callReadJSON(t, tools.listFindings(), `{}`).([]any) + if len(findings) != 1 { + t.Fatalf("list_findings: %#v", findings) + } + finding := findings[0].(map[string]any) + if finding["inherited"] != true || int64(finding["task_id"].(float64)) != source.ID || int64(finding["intent_id"].(float64)) != sourceIntent { + t.Fatalf("inherited finding provenance: %#v", finding) + } + + detailInput, _ := json.Marshal(map[string]any{"id": sourceFact}) + detail := callReadJSON(t, tools.nodeDetail(), string(detailInput)).(map[string]any) + if detail["inherited"] != true || int64(detail["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited node detail: %#v", detail) + } + + traceInput, _ := json.Marshal(map[string]any{"intent_id": sourceIntent}) + trace := callReadJSON(t, tools.getWorkerTrace(), string(traceInput)).(map[string]any) + if trace["inherited"] != true || int64(trace["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited trace: %#v", trace) + } + steps := trace["steps"].([]any) + if len(steps) != 1 || int64(steps[0].(map[string]any)["step_id"].(float64)) != stepID { + t.Fatalf("inherited trace steps: %#v", steps) + } + output := callReadJSON(t, tools.getWorkerOutput(), string(traceInput)).(map[string]any) + if output["inherited"] != true || output["final_text"] != "shared trace full detail" { + t.Fatalf("inherited worker output: %#v", output) + } + + search := callReadJSON(t, tools.searchAllWorkerTraces(), `{"q":"shared trace marker"}`).(map[string]any) + hits := search["hits"].([]any) + if len(hits) != 1 || hits[0].(map[string]any)["inherited"] != true { + t.Fatalf("inherited trace search: %#v", hits) + } +} diff --git a/agent/cancelcause.go b/agent/cancelcause.go new file mode 100644 index 0000000..f4705bc --- /dev/null +++ b/agent/cancelcause.go @@ -0,0 +1,90 @@ +package agent + +import ( + "context" + "errors" + "fmt" +) + +// AbortCause names why an agent run's context was cancelled. Every cancellation +// site should attach one so the activity trace can report the real initiator. +type AbortCause struct { + Code string + Short string + Text string +} + +func (c *AbortCause) Error() string { return c.Text } + +func cause(code, short, text string) *AbortCause { + return &AbortCause{Code: code, Short: short, Text: text} +} + +// Causef builds a cause that includes runtime-specific detail. +func Causef(code, short, format string, args ...any) *AbortCause { + return &AbortCause{Code: code, Short: short, Text: fmt.Sprintf(format, args...)} +} + +var ( + // Task-level execution context. + AbortPausedByUser = cause("paused_by_user", "用户暂停了任务", + "用户通过任务控制接口（POST /api/tasks/{id}/control，action=pause）暂停了任务。本次 Planner/Worker 运行被主动取消；运行中的意图会退回 frontier(open)，恢复任务后重新领取并从头执行") + AbortPausedByOrchestrator = cause("paused_by_orchestrator", "编排 Agent 暂停了任务", + "编排 Agent 调用了 pause_task 工具暂停本任务。本次 Planner/Worker 运行被主动取消；运行中的意图会退回 frontier(open)，恢复后重新执行") + AbortTaskDeleted = cause("task_deleted", "任务被删除", + "任务正在删除（DELETE /api/tasks/{id}），删除屏障已取消该任务正在运行的 Planner、Worker 和主 Agent；本次运行结果不会再被使用") + AbortPausedOnReload = cause("paused_on_reload", "后端恢复了任务的暂停状态", + "后端启动时根据数据库中持久化的状态恢复了任务暂停。本次运行被取消；正常情况下恢复阶段没有正在运行的 Agent") + AbortGoalMet = cause("goal_met", "规划者判定任务目标已达成", + "规划者判定任务目标已达成并将任务置为 done，随后取消仍在运行的 Worker；这些意图会标记为 stopped，而不是失败") + AbortSettleDrainTimeout = cause("settle_drain_timeout", "任务超时收尾的等待时间已用尽", + "任务到达 timeout 后等待正在运行的 Worker 优雅收尾，但 90 秒 drain 宽限仍不足，因此执行硬取消；意图会标记为 exhausted，收尾阶段已经写入的事实和资产会保留") + + // Per-work context. + AbortKilledByPlanner = cause("killed_by_planner", "规划者终止了这条意图", + "规划者调用 kill_work 主动终止了这条意图，通常表示方向跑偏或已无继续价值；意图会标记为 stopped，不会自动重新领取") + AbortWorkPausedByUser = cause("work_paused_by_user", "用户暂停了这条 Worker 意图", + "用户暂停了正在运行的 Worker。本次调用被取消，意图转为 paused；已经登记的意图、事实、漏洞和活动记录全部保留，恢复后从头重新执行") + AbortWorkCancelledByUser = cause("work_cancelled_by_user", "用户取消了这条 Worker 意图", + "用户取消了正在运行的 Worker。本次调用被取消；Worker 退出写入区后，服务端会事务性删除该意图及其直接产生的事实、漏洞和执行记录") + AbortWorkFinished = cause("work_finished", "Worker 已正常结束并释放 context", + "Worker 已正常结束，引擎在 detachWork 中释放其 context 资源。这不是运行中断；若它出现在中断消息中，说明取消与收场事件发生了竞态") + AbortPausedRaceGuard = cause("paused_race_guard", "任务暂停期间拒绝启动新运行", + "任务处于暂停状态时，引擎拒绝发出新的执行 context，用于防止 claim 与暂停之间的竞态导致 Worker 继续启动；已领取的意图会退回 frontier") + + // Main Agent and standalone conversation contexts. + AbortChatStoppedByUser = cause("chat_stopped_by_user", "用户停止了本轮对话", + "用户点击了停止，主动中止本轮主 Agent 或会话 Agent 运行。已经产生的活动记录会保留，可以继续发送下一条消息") + AbortChatPausedWithTask = cause("chat_paused_with_task", "任务暂停并中止了主 Agent 对话", + "用户暂停任务时，正在运行的主 Agent 对话也被同步取消。已经产生的活动记录会保留；恢复任务后不会自动重放本轮消息") + AbortChatTurnFinished = cause("chat_turn_finished", "本轮对话已正常结束并释放 context", + "本轮对话已正常结束，服务端正在释放该轮 context 资源。这不是运行中断；若它出现在中断消息中，说明取消与收场事件发生了竞态") + + // Process-level and per-run hard backstop. + AbortShutdown = cause("shutdown", "后端进程正在关闭", + "后端进程收到 SIGINT 或 SIGTERM，正在重启、更新或关闭。所有运行中的 Agent 会被取消；重启后残留的 running 意图会重置为 open 并重新执行") + AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底已触发", + "单次运行超过软墙钟预算及额外宽限，说明模型请求或某个工具长时间没有返回，导致正常的回合边界收尾无法执行。请重点检查中断前最后一个未返回的工具调用") +) + +// AbortReason resolves the named cause attached to a cancelled run context. +func AbortReason(ctx context.Context) (code, short, text string, ok bool) { + c := context.Cause(ctx) + if c == nil { + return "", "", "", false + } + var ac *AbortCause + if errors.As(c, &ac) { + return ac.Code, ac.Short, ac.Text, true + } + switch { + case errors.Is(c, context.DeadlineExceeded): + return "deadline_exceeded", "上游 context 到达 deadline", + "上游 context 到达 deadline，但设置方没有通过 WithTimeoutCause 附加具名原因: " + c.Error(), true + case errors.Is(c, context.Canceled): + return "canceled_no_cause", "取消方未附加具名原因", + "上游 context 被取消，但取消方没有通过 context.WithCancelCause 附加具名原因；请在 agent/cancelcause.go 登记原因并接入该取消点", true + default: + return "other", firstLine(c.Error(), 80), c.Error(), true + } +} diff --git a/agent/capture.go b/agent/capture.go index 419e855..965a9f3 100644 --- a/agent/capture.go +++ b/agent/capture.go @@ -3,11 +3,13 @@ package agent import ( "context" "strings" + "time" + "unicode/utf8" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/agentcore" "github.com/Autumn-27/norma/harness" "github.com/Autumn-27/norma/llm" - "github.com/Autumn-27/artex/db" ) // captureRun drives one agent turn-to-completion over Session.Prompt and emits a @@ -63,17 +65,20 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tkind = kind tbuf.WriteString(text) } + lastTool := &runTrace{startedAt: time.Now()} + var lastUsage *llm.Usage var finalText string var rerr error for ev, err := range s.Prompt(ctx, input) { if err != nil { flush() - if ctx.Err() != nil { // ctx cancelled = manual stop, not a failure - rec(db.Activity{Kind: "result", Summary: "（已手动停止本次运行）"}) + if ctx.Err() != nil { // engine/user cancellation, not a provider failure + sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool) + rec(activityWithUsage(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}, lastUsage)) return finalText, reason, ctx.Err() } - rec(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()}) + rec(activityWithUsage(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()}, lastUsage)) return finalText, reason, err } switch ev.Kind { @@ -84,6 +89,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, flush() toolNames[ev.ToolUse.ID] = ev.ToolUse.Name in := string(ev.ToolUse.Input) + lastTool.start(ev.ToolUse.ID, ev.ToolUse.Name, in) rec(db.Activity{Kind: "tool_use", Tool: ev.ToolUse.Name, ToolUseID: ev.ToolUse.ID, Summary: ev.ToolUse.Name + " " + firstLine(in, 200), Detail: in}) case harness.KindToolResult: @@ -92,6 +98,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, } flush() out := blocksText(ev.ToolResult.Content) + lastTool.done(ev.ToolResult.ToolUseID) rec(db.Activity{Kind: "tool_result", Tool: toolNames[ev.ToolResult.ToolUseID], ToolUseID: ev.ToolResult.ToolUseID, IsError: ev.ToolResult.IsError, Summary: firstLine(out, 200), Detail: out}) case harness.KindText: @@ -105,6 +112,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, // buffered final-answer text must stay for the KindResult de-dup. if ev.Usage != nil { u := *ev.Usage + lastUsage = &u rec(db.Activity{Kind: "usage", InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens, CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens}) @@ -120,25 +128,13 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tkind = "" } flush() // flush any trailing thinking / non-final text - sum := ev.Terminal.Text - if sum == "" { - switch ev.Terminal.Reason { - case harness.ReasonAbortedTools, harness.ReasonAbortedStreaming: - // interruption (task paused / worker killed / backend restart), - // not a real completion — the intent is re-queued (open) or - // stopped, so the session icon reflects that, not this run. - sum = "（运行被中断：任务暂停 / 终止 / 重启所致；未完成）" - case harness.ReasonTimeout, harness.ReasonMaxTurns: - // budget hit → settlement ran (facts/assets were written back); - // only the trailing text summary is missing, not a failure. - sum = "（达运行预算上限，已收尾写回事实；本次无文字总结，终态: " + string(ev.Terminal.Reason) + "）" - default: - sum = "（无总结，终态: " + string(ev.Terminal.Reason) + "）" - } + sum, detail := ev.Terminal.Text, ev.Terminal.Text + if sum == "" || ev.Terminal.Reason == harness.ReasonAbortedTools || ev.Terminal.Reason == harness.ReasonAbortedStreaming { + sum, detail = terminalText(ctx, ev.Terminal, lastTool) } u := ev.Terminal.Usage // cumulative token usage for this session rec(db.Activity{Kind: "result", IsError: ev.Terminal.Err != nil, - Summary: firstLine(sum, 400), Detail: sum, + Summary: firstLine(sum, 400), Detail: detail, InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens, CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens}) if ev.Terminal.Err != nil { @@ -151,6 +147,18 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, return finalText, reason, rerr } +func activityWithUsage(activity db.Activity, usage *llm.Usage) db.Activity { + if usage == nil { + return activity + } + u := *usage + activity.InputTokens = &u.InputTokens + activity.OutputTokens = &u.OutputTokens + activity.CacheReadTokens = &u.CacheReadTokens + activity.CacheWriteTokens = &u.CacheWriteTokens + return activity +} + // blocksText concatenates the text of a tool-result's content blocks. func blocksText(blocks []llm.ContentBlock) string { var b strings.Builder @@ -165,14 +173,14 @@ func blocksText(blocks []llm.ContentBlock) string { return b.String() } -// firstLine returns a single-line, length-capped preview for the summary column. +// firstLine returns a single-line, rune-capped preview for the summary column. func firstLine(s string, max int) string { s = strings.TrimSpace(s) - if i := strings.IndexByte(s, '\n'); i >= 0 { - s = s[:i] + if before, _, found := strings.Cut(s, "\n"); found { + s = before } - if len(s) > max { - s = s[:max] + "…" + if utf8.RuneCountInString(s) > max { + s = string([]rune(s)[:max]) + "…" } return s } diff --git a/agent/chat.go b/agent/chat.go index 86deccc..41be93b 100644 --- a/agent/chat.go +++ b/agent/chat.go @@ -92,6 +92,7 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin // visible skills/MCP and lets the DB tools table filter/override. DefaultTools // have no tools-table rows, so they always pass through. base := actool.DefaultTools() + ctx = WithRunInfo(ctx, RunInfo{SessionID: sessionID}) tools, def, cleanup := AugmentTools(ctx, agentKey, base) defer cleanup() diff --git a/agent/goals.go b/agent/goals.go index 52f5a22..388515a 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -7,6 +7,7 @@ import ( "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/agentcore" + "github.com/Autumn-27/norma/llm" acperm "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" ) @@ -74,6 +75,11 @@ type GoalSpec struct { // no provider is configured or the call yields nothing — the caller then falls // back to a rule-based split so goal nodes always exist. // +// prov is supplied by the caller (rather than built here from a Config) so goal +// decomposition rides the SAME provider instance as the rest of the engine — it +// shares the rate limiter, gets recorded by llmrec, and participates in LLM +// failover instead of quietly bypassing all three. +// // desc is the task's free-text description (背景：靶标范围/flag 数量/交战说明等). // It is fed alongside the goal so the decomposer no longer splits blind — the // prompt still forbids inventing anything the two texts don't state. @@ -88,9 +94,18 @@ type GoalSpec struct { // straight into it (the same managed tool the main agent uses to add goals at // runtime). The returned specs are read back from the store so callers can emit // per-goal activity and detect the "LLM produced nothing" case for their fallback. -func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { - prov, err := c.NewProvider() - if err != nil { +func DecomposeGoals(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { + if prov == nil { + return nil + } + return DecomposeGoalsWithProvider(ctx, prov, dataDir, goalText, desc, as, ts, taskID, emit) +} + +// DecomposeGoalsWithProvider is the task-runtime variant used when a task has an +// ordered provider chain. It preserves the same tools and write behavior while +// letting the caller own provider selection/failover. +func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { + if prov == nil { return nil } // worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link diff --git a/agent/mainagent.go b/agent/mainagent.go index 3c7291c..be9e573 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -21,6 +21,7 @@ type MainAgent struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) proxyAddr string // recording proxy for WebFetch (empty = direct) proxyCACert string // recording proxy's CA cert path (HTTPS verify) @@ -33,6 +34,15 @@ func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns} } +func (m *MainAgent) SetCompactionWindowResolver(fn func() int) { m.windowFn = fn } + +func (m *MainAgent) compactionWindow() int { + if m.windowFn != nil { + return m.windowFn() + } + return m.window +} + // SetProxy points the main agent's WebFetch at the recording proxy plus the CA // cert it trusts to verify HTTPS through it (empty addr = direct). func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert } @@ -77,13 +87,6 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t tsx.SetAssetStore(as, as.Companies()) } tsx.SetTaskID(taskID) -<<<<<<< Updated upstream - tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) - tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running - tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标：…」触发 - // 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash） - base := append(tsx.MainAgentTools(), actool.DefaultTools()...) -======= tsx.SetCoverageEnabled(as == nil || as.CoverageEnabled(taskID)) tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running @@ -93,7 +96,6 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t // 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets（不入 prompt）。 base := append(tsx.DropCoverageTools(tsx.MainAgentTools()), actool.DefaultTools()...) ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)}) ->>>>>>> Stashed changes tools, def, cleanup := AugmentTools(ctx, "mainagent", base) tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参 defer cleanup() @@ -121,9 +123,9 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA WorkingDir: mainDir, // 本任务工作目录 /tasks/ ToolOutputDir: cmdOutDir(mainDir), - MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) - Compaction: compactionConfig(m.window), // long chats stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 + MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) + Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 // 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。 Settlement: wrapupSettlement("mainagent", nil), } diff --git a/agent/planner.go b/agent/planner.go index 4ea20de..dfacc91 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -24,6 +24,7 @@ type Planner struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) killWork func(intentID int64) error // engine callback to terminate a running work (nil = off) steerWork func(intentID int64, msg string) error // engine callback to steer a running work mid-run (nil = off) @@ -45,6 +46,15 @@ func NewPlanner(prov llm.Provider, model, workDir string, tx *transcript.Store, return &Planner{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, todos: map[int64]*actool.TodoStore{}} } +func (p *Planner) SetCompactionWindowResolver(fn func() int) { p.windowFn = fn } + +func (p *Planner) compactionWindow() int { + if p.windowFn != nil { + return p.windowFn() + } + return p.window +} + // SetProxy points the planner's WebFetch at the recording proxy plus the CA cert // it trusts to verify HTTPS through it (empty addr = direct). func (p *Planner) SetProxy(addr, caCert string) { p.proxyAddr, p.proxyCACert = addr, caCert } @@ -103,12 +113,6 @@ func renderPlannerTodos(items []actool.Todo) string { // TriggerEvent describes what concretely caused this planning round to fire, so // the planner looks first at the actual change instead of re-scanning the whole // overview. Kind: -<<<<<<< Updated upstream -// "done" — a worker finished intent IntentID (its output conclusion is fetched). -// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). -// "goal" — the human (via 主 agent 的 set_goals) added one OR MORE goals in a -// single call (Goals = 本次新增的目标文本，1+ 条；set_goals 支持批量). -======= // // "done" — a worker finished intent IntentID (its output conclusion is fetched). // "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). @@ -118,7 +122,6 @@ func renderPlannerTodos(items []actool.Todo) string { // "goal_edited" — the human edited a goal from 总览的目标管理 (OldGoal→NewGoal 文本). // "cancelled" — the human deleted intent IntentID (Detail = 删除原因). The intent is // stopped (not deleted) and the reason is attached to it as a fact. ->>>>>>> Stashed changes type TriggerEvent struct { Kind string IntentID int64 @@ -291,13 +294,9 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts tsx.SetOwnerNode(origin) // planner-side anchors default to the task root (origin fact) } // 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash） -<<<<<<< Updated upstream - base := append(tsx.PlannerTools(), actool.DefaultTools()...) -======= // 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets（不入 prompt）。 base := append(tsx.DropCoverageTools(tsx.PlannerTools()), actool.DefaultTools()...) ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)}) ->>>>>>> Stashed changes tools, def, cleanup := AugmentTools(ctx, "planner", base) defer cleanup() // 关键态势（刚完成的意图 + 预取的完整图）改放【本轮 user 输入】(见下方 input)，system @@ -348,7 +347,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts ToolOutputDir: cmdOutDir(taskDir), MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management) MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余 - Compaction: compactionConfig(p.window), + Compaction: compactionConfig(p.compactionWindow()), // 跨唤醒共享的规划待办：让串行链在多轮之间保留（session 是新的，store 不是）。 Todos: p.todoFor(ts.ID()), // 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、 @@ -374,7 +373,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts runCtx := ctx if maxDur > 0 { var cancel context.CancelFunc - runCtx, cancel = context.WithTimeout(ctx, maxDur+settleHardGrace) + runCtx, cancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout) defer cancel() } _, _, err = captureRun(runCtx, opts, input, diff --git a/agent/runinfo.go b/agent/runinfo.go new file mode 100644 index 0000000..f006b6a --- /dev/null +++ b/agent/runinfo.go @@ -0,0 +1,47 @@ +package agent + +import ( + "context" + + "github.com/Autumn-27/artex/db" +) + +// RunInfo identifies WHICH run a tool call belongs to. Tool assembly only receives +// (ctx, agentKey) — the task/exploration ids live in the caller's arguments, not the +// ctx — so anything wired at assembly time (currently the Skill ledger in +// server/assembly.go) has no way to attribute a call to a task. Each run attaches +// its own RunInfo before calling AugmentTools; the wiring closure reads it once and +// captures it, so per-run attribution stays correct without threading parameters +// through the tool layer. Same pattern as TaskClock (see taskclock.go). +// +// Zero value = attribution unknown; every consumer must treat it as optional. +type RunInfo struct { + TaskID int64 // task registry id; 0 for non-task runs (chat sessions) + ExplorationID int64 // exploration id; 0 when unknown + IntentID int64 // worker's intent node; 0 for planner/mainagent/chat + SessionID string // chat conversation id; empty for task runs +} + +// explorationID reads a store's exploration id, tolerating a nil store (planner and +// worker runs can be driven without one in tests). +func explorationID(ts *db.ExplorationStore) int64 { + if ts == nil { + return 0 + } + return ts.ID() +} + +type runInfoKey struct{} + +// WithRunInfo attaches run attribution to ctx. +func WithRunInfo(ctx context.Context, ri RunInfo) context.Context { + return context.WithValue(ctx, runInfoKey{}, ri) +} + +// RunInfoFrom reads the RunInfo (zero value if none attached). +func RunInfoFrom(ctx context.Context) RunInfo { + if v, ok := ctx.Value(runInfoKey{}).(RunInfo); ok { + return v + } + return RunInfo{} +} diff --git a/agent/terminalreason.go b/agent/terminalreason.go new file mode 100644 index 0000000..838cbd1 --- /dev/null +++ b/agent/terminalreason.go @@ -0,0 +1,161 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/Autumn-27/norma/harness" +) + +// runTrace retains the latest tool call so an interrupted run can identify the +// operation that was still in flight. +type runTrace struct { + startedAt time.Time + id string + name string + input string + at time.Time + pending bool +} + +func (t *runTrace) start(id, name, input string) { + t.id, t.name, t.input, t.at, t.pending = id, name, input, time.Now(), true +} + +func (t *runTrace) done(id string) { + if id == t.id { + t.pending = false + } +} + +var reasonHint = map[harness.TerminalReason]string{ + harness.ReasonCompleted: "模型正常结束了本轮，但没有留下文字总结；事实和资产以本轮工具调用记录为准", + harness.ReasonMaxTurns: "达到步数上限(MaxTurns)：SDK 已执行收尾并写回事实和资产，意图会标记为 exhausted，供规划者换方向继续，而不是作为失败处理", + harness.ReasonTimeout: "达到单次运行的软墙钟预算(MaxDuration)：SDK 已在回合边界收尾并写回事实和资产，意图会标记为 exhausted", + harness.ReasonModelError: "模型或 API 调用失败（网络、鉴权、限流、供应商 5xx 等）；引擎会按策略退避重试，重试用尽后意图标记为 blocked", + harness.ReasonBlockingLimit: "上下文长度达到硬上限，请求在发出前被拦截；应收窄意图粒度或压缩工具返回", + harness.ReasonPromptTooLong: "提示词过长且上下文压缩重试已经用尽，无法继续执行", + harness.ReasonImageError: "当前模型不支持本轮多模态内容；请切换支持视觉的模型或避免工具返回图片", + harness.ReasonStopHookPrevented: "Stop 钩子阻止本轮结束，随后未能继续；请检查任务 Guard 规则是否过严", + harness.ReasonHookStopped: "工具或钩子主动停止继续执行，例如越界目标或禁用命令；请检查最后一条 tool_result 的拦截说明", + harness.ReasonAbortedStreaming: "运行在模型输出流式生成阶段被取消", + harness.ReasonAbortedTools: "运行在工具执行阶段被取消", +} + +// terminalText renders a terminal event with no final text into a compact summary +// and a Markdown detail block. +func terminalText(ctx context.Context, term *harness.Terminal, tr *runTrace) (string, string) { + reason := term.Reason + aborted := reason == harness.ReasonAbortedStreaming || reason == harness.ReasonAbortedTools + // Prompt may return ctx.Err directly without a terminal event. Preserve the + // cancellation cause instead of falling back to an empty/unknown terminal reason. + if reason == "" && ctx.Err() != nil { + aborted = true + } + + var sum string + if aborted { + _, short, _, ok := AbortReason(ctx) + if !ok { + short = "未能取得取消原因" + } + stage := "执行过程中" + switch reason { + case harness.ReasonAbortedStreaming: + stage = "模型输出阶段" + case harness.ReasonAbortedTools: + stage = "工具执行阶段" + } + sum = "（运行被中断：" + short + "；停在" + stage + progressSuffix(term, tr) + "，未完成）" + } else if reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout { + sum = "（达到运行预算上限(" + string(reason) + ")，已收尾写回事实" + progressSuffix(term, tr) + "；本次无文字总结）" + } else { + hint := terminalReasonHint(reason) + sum = "（无文字总结，终态 " + terminalReasonLabel(reason) + "：" + firstLine(hint, 80) + "）" + } + + var b strings.Builder + b.WriteString(sum) + b.WriteString("\n\n") + displayReason := terminalReasonLabel(reason) + fmt.Fprintf(&b, "- **终态**: `%s` - %s\n", displayReason, terminalReasonHint(reason)) + if aborted { + code, _, why, ok := AbortReason(ctx) + if ok { + fmt.Fprintf(&b, "- **中断原因** (`%s`): %s\n", code, why) + } else { + b.WriteString("- **中断原因**: 无法取得；取消方可能没有通过 context.WithCancelCause 附加具名原因\n") + } + } + if term.Err != nil { + fmt.Fprintf(&b, "- **底层错误**: `%v`\n", term.Err) + } + if aborted && strings.TrimSpace(term.Text) != "" { + b.WriteString("- **取消前已生成的部分输出**:\n\n") + b.WriteString(term.Text) + b.WriteString("\n\n") + } + if term.Turns > 0 { + fmt.Fprintf(&b, "- **已执行**: %d 轮模型回合\n", term.Turns) + } + if !tr.startedAt.IsZero() { + fmt.Fprintf(&b, "- **本次运行耗时**: %s\n", roundDur(time.Since(tr.startedAt))) + } + if u := term.Usage; u.InputTokens+u.OutputTokens+u.CacheReadTokens+u.CacheWriteTokens > 0 { + fmt.Fprintf(&b, "- **累计 token**: 输入 %d / 输出 %d / 缓存读 %d / 缓存写 %d\n", + u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheWriteTokens) + } + if tr.name == "" { + b.WriteString("- **工具调用**: 本次运行还没有发出工具调用就结束了\n") + } else if tr.pending { + fmt.Fprintf(&b, "- **中断时正在执行的工具**: `%s`（已运行 %s，**未返回结果**）\n\n ```json\n %s\n ```\n", + tr.name, roundDur(time.Since(tr.at)), firstLine(tr.input, 300)) + } else { + fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`（已正常返回）\n", tr.name) + } + return sum, b.String() +} + +func terminalReasonLabel(reason harness.TerminalReason) string { + if reason == "" { + return "context_canceled" + } + return string(reason) +} + +func terminalReasonHint(reason harness.TerminalReason) string { + if hint := reasonHint[reason]; hint != "" { + return hint + } + if reason == "" { + return "运行的 context 已取消，但底层没有产生 Terminal 事件" + } + return "未知终态；harness 可能新增了 TerminalReason，请补充 reasonHint" +} + +func progressSuffix(term *harness.Terminal, tr *runTrace) string { + var parts []string + if term.Turns > 0 { + parts = append(parts, fmt.Sprintf("%d 轮", term.Turns)) + } + if !tr.startedAt.IsZero() { + parts = append(parts, roundDur(time.Since(tr.startedAt))) + } + if len(parts) == 0 { + return "" + } + return "，已运行 " + strings.Join(parts, " / ") +} + +func roundDur(d time.Duration) string { + switch { + case d < time.Minute: + return d.Round(100 * time.Millisecond).String() + case d < time.Hour: + return d.Round(time.Second).String() + default: + return d.Round(time.Minute).String() + } +} diff --git a/agent/terminalreason_test.go b/agent/terminalreason_test.go new file mode 100644 index 0000000..ff53fe4 --- /dev/null +++ b/agent/terminalreason_test.go @@ -0,0 +1,169 @@ +package agent + +import ( + "context" + "errors" + "strings" + "testing" + "time" + "unicode/utf8" + + "github.com/Autumn-27/norma/harness" + "github.com/Autumn-27/norma/llm" +) + +func TestReasonHintCoversEveryTerminalReason(t *testing.T) { + all := []harness.TerminalReason{ + harness.ReasonCompleted, harness.ReasonBlockingLimit, harness.ReasonImageError, + harness.ReasonModelError, harness.ReasonAbortedStreaming, harness.ReasonAbortedTools, + harness.ReasonPromptTooLong, harness.ReasonStopHookPrevented, harness.ReasonHookStopped, + harness.ReasonMaxTurns, harness.ReasonTimeout, + } + for _, reason := range all { + if strings.TrimSpace(reasonHint[reason]) == "" { + t.Errorf("terminal reason %q has no explanation", reason) + } + } +} + +func TestAbortCausePropagatesThroughRunContextChain(t *testing.T) { + type key struct{} + execCtx, cancelExec := context.WithCancelCause(context.Background()) + workCtx, cancelWork := context.WithCancelCause(execCtx) + defer cancelWork(nil) + valued := context.WithValue(workCtx, key{}, "task-1") + runCtx, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) + defer cancelRun() + + cancelExec(AbortPausedByUser) + code, _, text, ok := AbortReason(runCtx) + if !ok || code != "paused_by_user" { + t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok) + } + if !strings.Contains(text, "frontier") { + t.Fatalf("cause detail did not propagate: %q", text) + } +} + +func TestAbortReasonFallbacks(t *testing.T) { + if _, _, _, ok := AbortReason(context.Background()); ok { + t.Fatal("live context must not report an abort reason") + } + plain, cancel := context.WithCancel(context.Background()) + cancel() + if code, _, _, ok := AbortReason(plain); !ok || code != "canceled_no_cause" { + t.Fatalf("code=%q ok=%v, want canceled_no_cause", code, ok) + } + timed, cancelTimed := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancelTimed() + <-timed.Done() + if code, _, _, _ := AbortReason(timed); code != "deadline_exceeded" { + t.Fatalf("code=%q, want deadline_exceeded", code) + } +} + +func TestTerminalTextAbortedNamesCauseAndHangingTool(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortKilledByPlanner) + trace := &runTrace{startedAt: time.Now().Add(-90 * time.Second)} + trace.start("tu_1", "Bash", `{"command":"nmap -p- 10.0.0.1"}`) + term := &harness.Terminal{ + Reason: harness.ReasonAbortedTools, + Err: context.Canceled, + Turns: 7, + Usage: llm.Usage{InputTokens: 1200, OutputTokens: 340}, + } + summary, detail := terminalText(ctx, term, trace) + if !strings.Contains(summary, "规划者") || strings.Contains(summary, "\n") { + t.Fatalf("unexpected summary: %q", summary) + } + for _, want := range []string{"killed_by_planner", "aborted_tools", "7 轮", "Bash", "未返回结果", "1200"} { + if !strings.Contains(detail, want) { + t.Errorf("detail missing %q:\n%s", want, detail) + } + } +} + +func TestTerminalTextDirectContextCancellationKeepsCause(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortChatStoppedByUser) + summary, detail := terminalText(ctx, &harness.Terminal{Err: context.Canceled}, &runTrace{startedAt: time.Now()}) + if !strings.Contains(summary, "用户停止") || !strings.Contains(detail, "chat_stopped_by_user") { + t.Fatalf("direct context cancellation lost cause: %q\n%s", summary, detail) + } + if !strings.Contains(detail, "context_canceled") { + t.Fatalf("missing synthesized terminal label: %s", detail) + } +} + +func TestTerminalTextAbortedPreservesPartialOutput(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortPausedByUser) + summary, detail := terminalText(ctx, &harness.Terminal{ + Reason: harness.ReasonAbortedStreaming, + Text: "已经生成的半段回答", + }, &runTrace{startedAt: time.Now()}) + if !strings.Contains(summary, "用户暂停") { + t.Fatalf("abort summary lost cause: %q", summary) + } + if !strings.Contains(detail, "已经生成的半段回答") { + t.Fatalf("abort detail lost partial output: %s", detail) + } +} + +func TestTerminalTextCompletedToolNotBlamed(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(AbortShutdown) + trace := &runTrace{startedAt: time.Now()} + trace.start("tu_1", "Read", `{"path":"/etc/hosts"}`) + trace.done("tu_1") + _, detail := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonAbortedStreaming}, trace) + if strings.Contains(detail, "未返回结果") || !strings.Contains(detail, "已正常返回") { + t.Fatalf("completed tool was blamed:\n%s", detail) + } +} + +func TestTerminalTextNonAbortReasons(t *testing.T) { + trace := &runTrace{startedAt: time.Now()} + summary, _ := terminalText(context.Background(), &harness.Terminal{Reason: harness.ReasonMaxTurns}, trace) + if !strings.Contains(summary, "运行预算上限") || strings.Contains(summary, "中断") { + t.Fatalf("unexpected max_turns summary: %q", summary) + } + summary, detail := terminalText(context.Background(), &harness.Terminal{ + Reason: harness.ReasonModelError, Err: errors.New("429 rate limited"), + }, trace) + if !strings.Contains(summary, "model_error") || !strings.Contains(detail, "429 rate limited") { + t.Fatalf("model_error detail incomplete: %q\n%s", summary, detail) + } +} + +func TestAbortCausesAreWellFormed(t *testing.T) { + all := []*AbortCause{ + AbortPausedByUser, AbortPausedByOrchestrator, AbortTaskDeleted, AbortPausedOnReload, + AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkPausedByUser, + AbortWorkCancelledByUser, AbortWorkFinished, AbortPausedRaceGuard, + AbortChatStoppedByUser, AbortChatPausedWithTask, AbortChatTurnFinished, + AbortShutdown, AbortRunHardTimeout, + } + seen := map[string]bool{} + for _, abort := range all { + switch { + case abort.Code == "" || seen[abort.Code]: + t.Errorf("missing or duplicate code: %q", abort.Code) + case abort.Short == "" || utf8.RuneCountInString(abort.Short) > 40: + t.Errorf("%s has invalid short text: %q", abort.Code, abort.Short) + case utf8.RuneCountInString(abort.Text) <= utf8.RuneCountInString(abort.Short): + t.Errorf("%s detail must be longer than short text", abort.Code) + } + seen[abort.Code] = true + } +} + +func TestFirstLineCapsByRuneNotByte(t *testing.T) { + if got := firstLine(strings.Repeat("恢复", 10), 5); got != "恢复恢复恢…" { + t.Fatalf("firstLine split a Unicode character: %q", got) + } + if got := firstLine("头\n尾", 100); got != "头" { + t.Fatalf("firstLine did not stop at newline: %q", got) + } +} diff --git a/agent/tools.go b/agent/tools.go index cfe0c6a..a1830f2 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -23,6 +23,10 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[ var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "summary": p["summary"], "state": n.State} + if n.Inherited { + m["source_task_id"] = n.SourceTaskID + m["inherited"] = true + } // asset_ids is the structured "which assets this direction covers" signal for // dedup; fall back to legacy payload keys (target_ids plural, then target_id // single) so intents stored before the rename still surface their anchors. @@ -355,34 +359,56 @@ func (t *ToolSet) graphOverview() actool.CoreTool { // the model needn't spend a turn calling the tool — every plan round starts with // an empty context and always needs this first). func (t *ToolSet) graphOverviewData() map[string]any { - { - out := map[string]any{} - // goals summary folded in so the planner needn't call list_goals each round. - goals, _ := t.ts.ListByKind(db.KindGoal, 100) - gsum := make([]map[string]any, 0, len(goals)) - for _, g := range goals { - var p map[string]any - _ = json.Unmarshal(g.Payload, &p) - gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + out := map[string]any{} + // goals summary folded in so the planner needn't call list_goals each round. + goals, _ := t.ts.ListByKind(db.KindGoal, 100) + gsum := make([]map[string]any, 0, len(goals)) + for _, g := range goals { + var p map[string]any + _ = json.Unmarshal(g.Payload, &p) + gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + } + out["goals"] = gsum + // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示；folded in so the + // planner reads them every round when generating intents (否则只写不读). + hints, _ := t.ts.ListByKind(db.KindHint, 50) + hsum := make([]map[string]any, 0, len(hints)) + for _, h := range hints { + var p map[string]any + _ = json.Unmarshal(h.Payload, &p) + hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + out["hints"] = hsum + // lineage from the exploration edges: an intent's parents (what it + // derived_from — possibly several facts combined) and its yields (the + // facts/findings it produced). factFrom maps a fact → the intent that + // produced it. This is the relationship layer the flat lists lacked. + edges, _ := t.ts.Edges(5000) + parentsOf := map[int64][]int64{} + yieldsOf := map[int64][]int64{} + factFrom := map[int64]int64{} + for _, e := range edges { + switch e.Rel { + case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent) + parentsOf[e.To] = append(parentsOf[e.To], e.From) + case db.RelYields: // intent --yields--> fact/finding + yieldsOf[e.From] = append(yieldsOf[e.From], e.To) + factFrom[e.To] = e.From } - out["goals"] = gsum - // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示；folded in so the - // planner reads them every round when generating intents (否则只写不读). - hints, _ := t.ts.ListByKind(db.KindHint, 50) - hsum := make([]map[string]any, 0, len(hints)) - for _, h := range hints { - var p map[string]any - _ = json.Unmarshal(h.Payload, &p) - hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + fr, _ := t.ts.Frontier(100) + out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) + all, _ := t.ts.ListByKind(db.KindIntent, 300) + var running, recentDone []*db.Node + for _, n := range all { + switch n.State { + case "running": + running = append(running, n) + case "done", "blocked", "exhausted": + if len(recentDone) < 15 { + recentDone = append(recentDone, n) + } } -<<<<<<< Updated upstream - out["hints"] = hsum - // lineage from the exploration edges: an intent's parents (what it - // derived_from — possibly several facts combined) and its yields (the - // facts/findings it produced). factFrom maps a fact → the intent that - // produced it. This is the relationship layer the flat lists lacked. - edges, _ := t.ts.Edges(5000) -======= } out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) @@ -460,6 +486,8 @@ func (t *ToolSet) graphOverviewData() map[string]any { e["value"] = r.Domain case r.Net != "": e["value"] = r.Net + case r.Value != "": + e["value"] = r.Value case r.CompanyID != nil: e["company_id"] = *r.CompanyID if t.cs != nil { @@ -654,107 +682,163 @@ func (t *ToolSet) relatedTaskOverviews() []map[string]any { stats, statsErr := ts.Stats() edges, _ := ts.Edges(5000) ->>>>>>> Stashed changes parentsOf := map[int64][]int64{} yieldsOf := map[int64][]int64{} factFrom := map[int64]int64{} - for _, e := range edges { - switch e.Rel { - case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent) - parentsOf[e.To] = append(parentsOf[e.To], e.From) - case db.RelYields: // intent --yields--> fact/finding - yieldsOf[e.From] = append(yieldsOf[e.From], e.To) - factFrom[e.To] = e.From + for _, edge := range edges { + switch edge.Rel { + case db.RelDerivedFrom, db.RelSpawns: + parentsOf[edge.To] = append(parentsOf[edge.To], edge.From) + case db.RelYields: + yieldsOf[edge.From] = append(yieldsOf[edge.From], edge.To) + factFrom[edge.To] = edge.From } } - fr, _ := t.ts.Frontier(100) - out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) - all, _ := t.ts.ListByKind(db.KindIntent, 300) - var running, recentDone []*db.Node - for _, n := range all { - switch n.State { - case "running": - running = append(running, n) - case "done", "blocked", "exhausted": - if len(recentDone) < 15 { - recentDone = append(recentDone, n) - } + + goals, _ := ts.ListByKind(db.KindGoal, relatedOverviewMaxGoalsPerSource) + goalSummary := make([]map[string]any, 0, len(goals)) + for _, goal := range goals { + var payload map[string]any + _ = json.Unmarshal(goal.Payload, &payload) + goalSummary = append(goalSummary, inheritedMap(map[string]any{ + "id": goal.ID, "state": goal.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["goals"] = goalSummary + + hints, _ := ts.ListByKind(db.KindHint, relatedOverviewMaxHintsPerSource) + hintSummary := make([]map[string]any, 0, len(hints)) + for _, hint := range hints { + var payload map[string]any + _ = json.Unmarshal(hint.Payload, &payload) + hintSummary = append(hintSummary, inheritedMap(map[string]any{ + "id": hint.ID, "state": hint.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["hints"] = hintSummary + + facts, _ := ts.ListByKind(db.KindFact, relatedOverviewMaxFactsPerSource) + findings, _ := ts.ListByKind(db.KindFinding, relatedOverviewMaxFindingsPerTask) + intentNodes, _ := ts.ListByKind(db.KindIntent, 300) + terminalIntent := make(map[int64]bool, len(intentNodes)) + for _, intent := range intentNodes { + terminalIntent[intent.ID] = inheritedIntentSummaryState(intent.State) + } + item["facts"] = len(facts) + item["findings"] = len(findings) + if statsErr == nil { + item["facts"] = stats[db.KindFact] + item["findings"] = stats[db.KindFinding] + if stats[db.KindGoal] > len(goals) || stats[db.KindHint] > len(hints) || + stats[db.KindFact] > len(facts) || stats[db.KindFinding] > len(findings) { + budget.truncated = true } } - out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) - out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) - out["frontier_open"] = len(fr) - // findings (confirmed vulns) and facts (worker exploration results) are - // now distinct node kinds. recent_facts surfaces fact summaries (esp. - // negative results) so the planner sees them in one call; full content - // via node_detail(id). - vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000) - factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first - out["findings"] = len(vulnNodes) // 确认漏洞数（目标判定看它） - out["facts"] = len(factNodes) // 探索事实/结论数（含否定结论） - recentFacts := make([]map[string]any, 0, 20) - for _, n := range factNodes { - if len(recentFacts) >= 20 { - break + recentFindings := make([]map[string]any, 0, len(findings)) + for _, finding := range findings { + entry := inheritedMap(compactFinding(finding), source.Task.TaskID) + entry["summary"] = budget.take(entry["summary"], 400) + recentFindings = append(recentFindings, entry) + } + item["recent_findings"] = recentFindings + recentFacts := make([]map[string]any, 0, len(facts)) + for _, fact := range facts { + m := inheritedMap(compactNode(fact), source.Task.TaskID) + m["summary"] = budget.take(m["summary"], 400) + if from := factFrom[fact.ID]; from > 0 && terminalIntent[from] { + m["from_intent"] = from } - m := compactNode(n) - if from := factFrom[n.ID]; from > 0 { - m["from_intent"] = from // 本事实由哪个意图产生 - } - // confidence 带进概览：让规划者一眼看出哪条结论只是 inferred（尤其否定结论 - // 别当铁案）；evidence 较长，留给 node_detail(id)。 - var fp map[string]any - if json.Unmarshal(n.Payload, &fp) == nil { - if c, ok := fp["confidence"].(string); ok && c != "" { - m["confidence"] = c + var payload map[string]any + if json.Unmarshal(fact.Payload, &payload) == nil { + if confidence, ok := payload["confidence"].(string); ok && confidence != "" { + m["confidence"] = confidence } } recentFacts = append(recentFacts, m) } - out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?}，详情/证据用 node_detail(id) - // the original task (root) so the planner always has it, not just the - // decomposed goals. - if description, goal, err := t.ts.Root(); err == nil { - out["task"] = map[string]any{"description": description, "goal": goal} + item["recent_facts"] = recentFacts + + recentDone := recentTerminalIntents(ts, relatedOverviewMaxIntentsPerTask) + for _, intent := range recentDone { + intent.Inherited = true + intent.SourceTaskID = source.Task.TaskID } - // coverage：粗略的资产测试覆盖度参考——范围(task_scope)内的资产里，被 fact 碰过的 - // 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。 - if t.as != nil && t.ts != nil && t.taskID > 0 { - if cov, err := t.as.TaskCoverage(t.taskID, t.ts.ID()); err == nil { - m := map[string]any{ - "denominator": cov.Denominator, - "tested": cov.Tested, - "by_type": cov.ByType, - "note": "coverage资产测试覆盖度（包括接口等各种相关资产），粗略估计、仅供参考：容器型资产/大量枚举会让它偏低，勿据此认为已测完；scope 是当前任务测试范围的根资产（root_domain/subdomain/ip/cidr/company 原始行），即覆盖度分母的边界、本任务授权/圈定的目标本身；可用 add_task_scope 增补范围、list_untested_assets 看未测的具体资产【通常不调用list_untested_assets，按照任务推进即可】；", + intentResults := compactIntents(recentDone, parentsOf, yieldsOf) + for i, intent := range recentDone { + intentResults[i]["summary"] = budget.take(intentResults[i]["summary"], 400) + acts, _, err := ts.ActivityPageForTerminalIntent(intent.ID, 0, 20) + if err != nil { + continue + } + var resultSummary, textFallback string + for _, activity := range acts { + switch activity.Kind { + case "result": + resultSummary = activity.Summary + case "text": + textFallback = activity.Summary } - if cov.ScopeRows == 0 { - m["pct"] = nil - m["status"] = "范围未锚定" - } else { - m["pct"] = cov.Pct - } - // scope：当前测试范围的根资产（task_scope 原始行），让 agent 知道这个任务到底 - // 圈定了哪些目标（不是全部测试资产，而是范围边界本身）。 - if rows, err := t.as.ListTaskScope(t.taskID); err == nil && len(rows) > 0 { - scope := make([]map[string]any, 0, len(rows)) - for _, r := range rows { - e := map[string]any{"kind": r.Kind, "source": r.Source} - switch { - case r.Domain != "": - e["value"] = r.Domain - case r.Net != "": - e["value"] = r.Net - case r.CompanyID != nil: - e["company_id"] = *r.CompanyID - } - scope = append(scope, e) - } - m["scope"] = scope - } - out["coverage"] = m + } + if resultSummary == "" { + resultSummary = textFallback + } + if resultSummary != "" { + intentResults[i]["result_summary"] = budget.take(resultSummary, 800) } } - return out + item["recent_intent_results"] = intentResults + if statsErr == nil { + item["node_stats"] = stats + } + + if t.as != nil { + if scopeRows, err := t.as.ListTaskScope(source.Task.TaskID); err == nil && len(scopeRows) > 0 { + scopeCount := len(scopeRows) + if len(scopeRows) > relatedOverviewMaxScopePerSource { + scopeRows = scopeRows[:relatedOverviewMaxScopePerSource] + budget.truncated = true + } + scope := make([]map[string]any, 0, len(scopeRows)) + for _, row := range scopeRows { + entry := map[string]any{"kind": row.Kind, "source": budget.take(row.Source, 300)} + switch { + case row.Domain != "": + entry["value"] = budget.take(row.Domain, 400) + case row.Net != "": + entry["value"] = budget.take(row.Net, 400) + case row.Value != "": + entry["value"] = budget.take(row.Value, 400) + case row.CompanyID != nil: + entry["company_id"] = *row.CompanyID + } + scope = append(scope, entry) + } + item["asset_scope"] = scope + item["asset_scope_count"] = scopeCount + } + if coverage, err := t.as.TaskCoverage(source.Task.TaskID, source.Task.ExplorationID); err == nil { + item["asset_coverage"] = map[string]any{ + "denominator": coverage.Denominator, + "tested": coverage.Tested, + "pct": coverage.Pct, + "by_type": coverage.ByType, + } + } + } + if budget.truncated { + item["summary_truncated"] = true + } + out = append(out, item) + } + return out +} + +func inheritedIntentSummaryState(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false } } @@ -763,7 +847,11 @@ func (t *ToolSet) relatedTaskOverviews() []map[string]any { func compactNode(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) - return map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } + return m } // compactFinding is compactNode plus the vuln-specific vulnclass/severity. @@ -771,6 +859,9 @@ func compactFinding(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } if vc, ok := p["vulnclass"]; ok && vc != nil && vc != "" { m["vulnclass"] = vc } @@ -781,16 +872,23 @@ func compactFinding(n *db.Node) map[string]any { } func (t *ToolSet) listFindings() actool.CoreTool { - return readTool("list_findings", "列【确认漏洞】(紧凑：id+task_id+intent_id+vulnclass+severity+摘要+状态)。task_id=该漏洞所属任务, intent_id=产生它的意图。这里只有漏洞,不含普通探索事实(那是 list_facts)。详情用 node_detail(id)。", + return readTool("list_findings", "列本任务及直接关联任务的【确认漏洞】(紧凑：id+task_id+intent_id+vulnclass+severity+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。这里只含漏洞；普通探索事实用 list_facts，详情用 node_detail(id)。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFinding, 500) - intentOf, _ := t.ts.FindingIntents() // finding id -> 产生它的 intent id - taskID := t.ts.ID() // exploration id = 任务 id（本 store 内所有 finding 同属） + f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500) + intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id + taskID := t.taskID + if taskID <= 0 { + taskID, _ = t.ts.TaskID() + } out := make([]map[string]any, 0, len(f)) for _, n := range f { m := compactFinding(n) - m["task_id"] = taskID + if n.Inherited { + m["task_id"] = n.SourceTaskID + } else { + m["task_id"] = taskID + } if iid, ok := intentOf[n.ID]; ok { m["intent_id"] = iid } @@ -801,10 +899,10 @@ func (t *ToolSet) listFindings() actool.CoreTool { } func (t *ToolSet) listFacts() actool.CoreTool { - return readTool("list_facts", "列【探索事实/结论】(紧凑：id+摘要+状态),即 worker 按意图探出的结论(含'端口关闭/不可注入'等否定结论)。详情用 node_detail(id)。漏洞看 list_findings。", + return readTool("list_facts", "列本任务及直接关联任务的【探索事实/结论】(紧凑：id+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。详情用 node_detail(id)，漏洞看 list_findings。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFact, 500) + f, _ := t.ts.ListByKindWithSources(db.KindFact, 500) out := make([]map[string]any, 0, len(f)) for _, n := range f { out = append(out, compactNode(n)) @@ -814,7 +912,7 @@ func (t *ToolSet) listFacts() actool.CoreTool { } func (t *ToolSet) nodeDetail() actool.CoreTool { - return readTool("node_detail", "按 id 取某个【探索图节点】的完整内容(事实/发现/意图/目标：摘要 + 详情/证据/PoC)。仅限探索图节点 id(来自 list_facts / list_findings / graph_overview 的 recent_facts 返回的 id)。查【资产】请用 list_assets / asset_neighbors——资产 id 不是探索节点 id,不要传进来。", + return readTool("node_detail", "按 id 取本任务或直接关联任务的【探索图节点】完整内容。继承节点带 source_task_id/inherited=true 且只读。仅限 list_facts/list_findings/graph_overview 返回的探索节点 id；资产请用 list_assets/asset_neighbors。", obj(map[string]any{"id": idp("探索图节点 id(非资产 id)")}, "id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -825,7 +923,7 @@ func (t *ToolSet) nodeDetail() actool.CoreTool { if id <= 0 { return actool.Errorf("id 必填"), nil } - n, err := t.ts.GetNode(id) + n, err := t.ts.GetNodeWithSources(id) if err != nil { return actool.Errorf(err.Error()), nil } @@ -857,9 +955,9 @@ func (t *ToolSet) addOneIntent(it intentItem) (int64, error) { // 先校验锚点（建节点前，避免坏锚点留下孤儿意图）。 parents := pidList(it.ParentIDs) for _, pidv := range parents { - n, err := t.ts.GetNode(pidv) + n, err := t.ts.GetNodeWithSources(pidv) if err != nil || n == nil { - return 0, fmt.Errorf("parent_id %d 不存在：parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id；顶层全新方向请留空 parent_ids", pidv) + return 0, fmt.Errorf("parent_id %d 不存在于本任务或直接关联任务：parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id；顶层全新方向请留空 parent_ids", pidv) } if n.Kind != db.KindFact && n.Kind != db.KindFinding { return 0, fmt.Errorf("parent_id %d 是 %q 节点，不能作为意图锚点：意图只能锚在已确认的【事实(fact)/发现(finding)】上，不能挂在意图/目标/提示上；顶层全新方向请留空 parent_ids", pidv, n.Kind) @@ -968,6 +1066,14 @@ func (t *ToolSet) proveGoal() actool.CoreTool { if goal == 0 || ev == 0 { return actool.Errorf("goal_id 和 evidence_id 必填"), nil } + goalNode, err := t.ts.GetNode(goal) + if err != nil || goalNode == nil || goalNode.Kind != db.KindGoal { + return actool.Errorf("goal_id 必须是本任务的目标节点（关联任务目标只读）"), nil + } + evidenceNode, err := t.ts.GetNodeWithSources(ev) + if err != nil || evidenceNode == nil || (evidenceNode.Kind != db.KindFact && evidenceNode.Kind != db.KindFinding) { + return actool.Errorf("evidence_id 必须是本任务或直接关联任务的事实/漏洞节点"), nil + } _ = t.ts.Link(ev, db.RelProves, goal) _ = t.ts.SetNodeState(goal, "met") // 每标记一个目标 met，就检查本任务是否【所有目标】都已 met；若是，自动判定 @@ -1032,12 +1138,19 @@ func (t *ToolSet) addFinding() actool.CoreTool { } var id int64 if t.ts != nil { + intent := pid(a.IntentID) + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } + } var err error id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors) if err != nil { return actool.Errorf(err.Error()), nil } - if intent := pid(a.IntentID); intent > 0 { + if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding } _, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors) @@ -1092,15 +1205,21 @@ func (t *ToolSet) recordOneFact(it factItem, defaultIntent int64) (int64, error) if c := strings.TrimSpace(it.Confidence); c != "" { payload["confidence"] = c } + intent := pid(it.IntentID) + if intent <= 0 { + intent = defaultIntent + } + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return 0, fmt.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）") + } + } // a fact is its OWN node kind (distinct from a vuln finding). id, err := t.ts.AddNode(db.KindFact, payload, 5, "confirmed", t.worker, pidList(it.AssetIDs)) if err != nil { return 0, err } - intent := pid(it.IntentID) - if intent <= 0 { - intent = defaultIntent - } if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> fact } @@ -1422,6 +1541,10 @@ func (t *ToolSet) killWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } if err := t.killWork(id); err != nil { return actool.Errorf(err.Error()), nil } @@ -1452,6 +1575,10 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } if strings.TrimSpace(a.Message) == "" { return actool.Errorf("message 必填"), nil } @@ -1464,7 +1591,7 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { // getWorkerOutput returns a work's final (or截至中止时的) conclusion text by intent id. func (t *ToolSet) getWorkerOutput() actool.CoreTool { - return readTool("get_worker_output", "取某条意图(work)的最终输出结论。正常结束返回其总结；被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", + return readTool("get_worker_output", "取本任务或直接关联任务某条意图(work)的最终输出结论。关联任务结果带 source_task_id/inherited=true 且只读。正常结束返回其总结；被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", obj(map[string]any{"intent_id": idp("意图 id（= work 句柄）")}, "intent_id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -1475,7 +1602,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } - acts, _, err := t.ts.ActivityList(&id, 0, 1000) + intentNode, err := t.ts.GetNodeWithSources(id) + if err != nil { + return actool.Errorf(err.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } + acts, _, err := t.ts.ActivityListWithSources(id, 0, 1000) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1494,16 +1628,25 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { pick, terminated = fallback, true } if pick == nil { + if intentNode.Inherited { + return jsonResult(inheritedMap(map[string]any{ + "intent_id": id, "final_text": "（该 work 尚无任何输出）", "terminated": true, + }, intentNode.SourceTaskID)) + } return actool.Text("（该 work 尚无任何输出）"), nil } - detail, _ := t.ts.ActivityDetail(pick.ID) + detail, _ := t.ts.ActivityDetailWithSources(pick.ID) if detail == "" { detail = pick.Summary } - return jsonResult(map[string]any{ + result := map[string]any{ "intent_id": id, "worker_name": pick.Worker, "final_text": detail, "summary": pick.Summary, "is_error": pick.IsError, "terminated": terminated, - }) + } + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1513,10 +1656,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { func traceSteps(acts []db.Activity) []map[string]any { steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } return steps } @@ -1530,7 +1677,7 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { "① 只传 intent_id → 返回该 work 每一步的摘要流（summary≤100字，含 step_id；只是动作轮廓，不含完整输出）；\n"+ "② intent_id + q → 只返回命中关键字的步骤摘要（在摘要和完整输出里都搜；仍只给 summary，要看内容用③）；\n"+ "③ intent_id + step_ids（最多5个）→ 返回这些步骤的完整内容(detail)。\n"+ - "典型流程：先①/②定位可疑步骤的 step_id，再用③取其完整输出。不含思考(thinking)步骤。", + "典型流程：先①/②定位可疑步骤的 step_id，再用③取其完整输出。不含思考(thinking)步骤。支持直接关联任务的历史 trace；其结果带 source_task_id/inherited=true 且只读。", obj(map[string]any{ "intent_id": idp("意图 id（= work 句柄）"), "q": str("关键字：只返回摘要/完整输出命中它的步骤（可选；与 step_ids 互斥）"), @@ -1549,6 +1696,13 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + intentNode, nodeErr := t.ts.GetNodeWithSources(id) + if nodeErr != nil { + return actool.Errorf(nodeErr.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } // ③ detail drill-down by step ids (max 5), thinking excluded by the store. if len(a.StepIDs) > 0 { var ids []int64 @@ -1560,31 +1714,47 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if len(ids) > 5 { return actool.Errorf("step_ids 一次最多 5 个（先用不带 step_ids 的调用定位，再分批取）"), nil } - acts, err := t.ts.ActivityByIDs(ids) + acts, err := t.ts.ActivityByIDsWithSources(ids) if err != nil { return actool.Errorf(err.Error()), nil } steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + if acts[i].NodeID == nil || *acts[i].NodeID != id || acts[i].Inherited != intentNode.Inherited || + (acts[i].Inherited && acts[i].SourceTaskID != intentNode.SourceTaskID) { + continue + } + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "detail": acts[i].Detail, - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } - return jsonResult(map[string]any{"intent_id": id, "steps": steps}) + result := map[string]any{"intent_id": id, "steps": steps} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) } // ①/② summary stream, optionally keyword-filtered; 100-char summaries. var acts []db.Activity var err error if strings.TrimSpace(a.Q) != "" { - acts, err = t.ts.ActivityTraceSearch(&id, a.Q, a.Limit) + acts, err = t.ts.ActivityTraceSearchWithSources(id, a.Q, a.Limit) } else { - acts, err = t.ts.ActivityTrace(id, a.Limit) + acts, err = t.ts.ActivityTraceWithSources(id, a.Limit) } if err != nil { return actool.Errorf(err.Error()), nil } - return jsonResult(map[string]any{"intent_id": id, "steps": traceSteps(acts)}) + result := map[string]any{"intent_id": id, "steps": traceSteps(acts)} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1610,7 +1780,7 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { return actool.Errorf("q 必填"), nil } // 排除调用者自身这条意图的步骤（worker 的自有 trace 已在其上下文里）。 - acts, err := t.ts.ActivityTraceSearchExcluding(t.ownerNode, a.Q, a.Limit) + acts, err := t.ts.ActivityTraceSearchAllWithSources(t.ownerNode, a.Q, a.Limit) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1620,11 +1790,15 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { if acts[i].NodeID != nil { intent = *acts[i].NodeID } - hits = append(hits, map[string]any{ + hit := map[string]any{ "intent_id": intent, "step_id": acts[i].ID, "worker": acts[i].Worker, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(hit, acts[i].SourceTaskID) + } + hits = append(hits, hit) } return jsonResult(map[string]any{"query": a.Q, "hits": hits}) }) @@ -1655,13 +1829,16 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if limit <= 0 { limit = 50 } - all, err := t.ts.ListByKind(db.KindIntent, 500) + all, err := t.ts.ListByKindWithSources(db.KindIntent, 500) if err != nil { return actool.Errorf(err.Error()), nil } q := strings.ToLower(strings.TrimSpace(a.Q)) out := make([]map[string]any, 0, limit) for _, n := range all { + if n.Inherited && n.State == "running" { + continue + } switch n.State { case "running", "done", "exhausted", "blocked", "stopped": // has run → has a process default: @@ -1673,7 +1850,11 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if q != "" && !strings.Contains(strings.ToLower(summary), q) { continue } - out = append(out, map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State}) + item := map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State} + if n.Inherited { + inheritedMap(item, n.SourceTaskID) + } + out = append(out, item) if len(out) >= limit { break } diff --git a/agent/tools_insert.go b/agent/tools_insert.go index 8019b42..5b4d58c 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "fmt" "strconv" "strings" @@ -284,6 +285,18 @@ func (t *ToolSet) insertAssets() actool.CoreTool { results = append(results, result{Index: i, ID: id, Type: typ}) t.writes.Assets++ t.anchorOwner(id) + if taskID > 0 { + var sourceNodeID *int64 + if t.ownerNode > 0 { + nodeID := t.ownerNode + sourceNodeID = &nodeID + } + summary := "Agent 通过 insert_assets 登记" + if t.ownerNode > 0 { + summary = fmt.Sprintf("Worker 意图 #%d 通过 insert_assets 登记", t.ownerNode) + } + _ = t.as.SetTaskAssetSource(taskID, id, "agent", summary, sourceNodeID) + } // 自动入测试范围(source='auto')：只对 worker 顶层显式插入的这一项，按其 // 类型加保守范围；side-effect 派生的资产不经此处，故范围不盲目扩大。taskID=0 时无操作。 // 资产覆盖度功能关闭时不再累积测试范围(分母)；related=false(与当前任务无关) @@ -310,14 +323,14 @@ func (t *ToolSet) insertAssets() actool.CoreTool { func (t *ToolSet) addCompanyScope() actool.CoreTool { return writeTool( "add_company_scope", - "把根域名/IP/IP段(CIDR)加入某公司的【资产范围】——加入后系统自动认领所有命中的资产(现有+未来)。\n"+ + "把域名/IP/CIDR/ICP备案/企业关键词加入某公司的【资产范围】——域名、网络和ICP会自动认领命中的资产，关键词只提供给Agent作为范围提示。\n"+ "公司名唯一：company 不存在则新建，已存在则复用(只把范围并进去)。\n"+ - "scope 一行一条，支持：根域名 / 单个 IP / CIDR 网段(如 1.2.3.0/24)。\n"+ + "scope 一行一条，系统自动识别：根域名 / URL / 单个 IP / CIDR 网段 / ICP备案 / 企业关键词。\n"+ "务必给 reason 说明归属依据(whois/证书/ASN 等)。\n"+ - "护栏：拒绝裸 TLD 与过宽网段(IPv4<=/16、IPv6<=/32)，非法行会被跳过并在 errors 返回。", + "护栏：拒绝裸 TLD 与过宽网段(IPv4前缀需为/16-/32、IPv6前缀需为/32-/128)，非法行会被跳过并在 errors 返回。", obj(map[string]any{ "company": str("公司名(不存在则新建、存在则复用；名称唯一)"), - "scope": str("资产范围，一行一条：根域名 / IP / CIDR"), + "scope": str("资产范围，一行一条：域名 / URL / IP / CIDR / ICP备案 / 企业关键词"), "reason": str("归属依据(证据/来源)，务必填写"), "logo": str("公司图标 URL(可选；仅新建公司时生效)"), }, "company", "scope"), @@ -365,14 +378,14 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { return writeTool( "add_task_scope", "把测试范围加入【本任务】——这是资产测试覆盖度的分母，也是本任务的授权边界。\n"+ - "kind 支持：company(整个公司名下资产) / root_domain(整个根域，含所有子域) / subdomain(单个精确子域) / ip / cidr。\n"+ + "kind 支持：company(整个公司名下资产) / root_domain(整个根域，含所有子域) / subdomain(单个精确子域) / ip / cidr / icp / keyword。\n"+ "说明：worker 逐个碰到的主机会被系统【自动】加进范围(精确子域)；本工具用于【主动扩大】——把整个根域/整个公司纳入，或补充指定某子域/IP。\n"+ - "value：company 传公司名或 id(公司须已存在)；root_domain/subdomain 传域名；ip/cidr 传 IP 或网段。\n"+ + "value：company 传公司名或 id(公司须已存在)；root_domain/subdomain 传域名；ip/cidr 传 IP 或网段；icp/keyword 传备案号或企业关键词。\n"+ "务必给 reason 说明依据(可审计)。多条用 entries 数组。", obj(map[string]any{ - "entries": map[string]any{"type": "array", "description": "批量：[{kind, value}]。kind∈company/root_domain/subdomain/ip/cidr；value=公司名或id / 域名 / IP / CIDR。", "items": map[string]any{"type": "object"}}, - "kind": str("[单条] company / root_domain / subdomain / ip / cidr"), - "value": str("[单条] 公司名或id / 域名 / IP / CIDR"), + "entries": map[string]any{"type": "array", "description": "批量：[{kind, value}]。kind∈company/root_domain/subdomain/ip/cidr/icp/keyword。", "items": map[string]any{"type": "object"}}, + "kind": str("[单条] company / root_domain / subdomain / ip / cidr / icp / keyword"), + "value": str("[单条] 公司名或id / 域名 / IP / CIDR / ICP / 关键词"), "reason": str("加入依据(用于审计)，务必填写"), }), func(_ context.Context, in json.RawMessage) (actool.Result, error) { @@ -399,12 +412,12 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { var added []map[string]any errs := map[string]string{} for i, e := range items { - ts, err := t.as.AddAgentScope(t.taskID, strings.TrimSpace(e.Kind), e.Value, a.Reason) + ts, err := t.as.AddAgentScope(t.taskID, strings.TrimSpace(e.Kind), e.Value, a.Reason, "agent") if err != nil { errs[strconv.Itoa(i)] = err.Error() continue } - added = append(added, map[string]any{"kind": ts.Kind, "domain": ts.Domain, "net": ts.Net, "company_id": ts.CompanyID}) + added = append(added, map[string]any{"kind": ts.Kind, "domain": ts.Domain, "net": ts.Net, "value": ts.Value, "company_id": ts.CompanyID}) } out := map[string]any{"added": added} if len(errs) > 0 { @@ -415,13 +428,12 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { ) } -// listUntestedAssets lets the plan agent pull the CURRENT TASK's in-scope, not-yet- -// tested assets on demand (filter by type, paginated) — so it can decide what to -// test next itself, instead of coverage pushing a backlog list into every prompt. +// listUntestedAssets lets the plan agent pull the current + directly inherited +// scope's not-yet-tested assets on demand (filter by type, paginated). func (t *ToolSet) listUntestedAssets() actool.CoreTool { return readTool( "list_untested_assets", - "查询【本任务】范围内、还没被测过的资产（供你自己判断要不要补测，不代替你决策）。\n"+ + "查询【本任务及直接关联任务】范围内、还没被事实锚点覆盖的资产（关联范围只读，供你自己判断要不要补测，不代替你决策）。\n"+ "可选按资产类型过滤：root_domain/subdomain/service/app/endpoint/ip。\n"+ "分页：page 从 1 起、page_size 默认 10。返回 {assets:[{id,type,label}], total, page, page_size}。仅任务上下文可用。", obj(map[string]any{ @@ -449,7 +461,7 @@ func (t *ToolSet) listUntestedAssets() actool.CoreTool { a.PageSize = 10 } offset := (a.Page - 1) * a.PageSize - assets, total, err := t.as.ListUntestedAssets(t.taskID, t.ts.ID(), strings.TrimSpace(a.Type), a.PageSize, offset) + assets, total, err := t.as.ListUntestedAssetsWithSources(t.taskID, strings.TrimSpace(a.Type), a.PageSize, offset) if err != nil { return actool.Errorf(err.Error()), nil } diff --git a/agent/tools_overview_test.go b/agent/tools_overview_test.go new file mode 100644 index 0000000..8fa6a80 --- /dev/null +++ b/agent/tools_overview_test.go @@ -0,0 +1,31 @@ +package agent + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/Autumn-27/artex/db" +) + +func TestOverviewTextBudgetIsFairAndUTF8Safe(t *testing.T) { + if got := relatedOverviewBudgetForSources(1); got != relatedOverviewMaxTextPerSource { + t.Fatalf("single source budget=%d want=%d", got, relatedOverviewMaxTextPerSource) + } + perSource := relatedOverviewBudgetForSources(db.MaxTaskSourceCount) + if perSource*db.MaxTaskSourceCount > relatedOverviewTotalTextRunes { + t.Fatalf("aggregate budget exceeded: per_source=%d", perSource) + } + + budget := overviewTextBudget{remaining: 5} + got := budget.take(strings.Repeat("中", 10), 20) + if !utf8.ValidString(got) { + t.Fatalf("budget truncation produced invalid UTF-8: %q", got) + } + if utf8.RuneCountInString(got) != 5 || !budget.truncated || budget.remaining != 0 { + t.Fatalf("unexpected truncation: got=%q budget=%+v", got, budget) + } + if tail := budget.take("more", 20); tail != "" { + t.Fatalf("exhausted budget returned more text: %q", tail) + } +} diff --git a/agent/worker.go b/agent/worker.go index 0de83ec..f2e5ce8 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -49,6 +49,7 @@ type Worker struct { mem *memory.Store // cross-engagement tradecraft memory (G4) tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) // runTimeout is the wall-clock budget for the main exploration of one intent // (0 = unlimited). When it fires, the run is cut and a settlement round is @@ -95,6 +96,15 @@ func NewWorker(prov llm.Provider, model, workDir string, tx *transcript.Store, w return &Worker{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, extraTools: extra} } +func (w *Worker) SetCompactionWindowResolver(fn func() int) { w.windowFn = fn } + +func (w *Worker) compactionWindow() int { + if w.windowFn != nil { + return w.windowFn() + } + return w.window +} + // SetProxy configures the recording proxy address that workers route target // traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through // that MITM proxy. Empty addr disables the hint. @@ -268,6 +278,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. // phase, Bash is hidden via Settlement.DisabledTools (no local gating needed). base := append(tsx.WorkerTools(), w.extraTools...) base = append(base, actool.DefaultTools()...) + ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts), IntentID: intent.ID}) tools, def, cleanup := AugmentTools(ctx, "worker", base) tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参 defer cleanup() @@ -324,8 +335,8 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. // large tool output spills to cmd-output/ with a head + pointer (SDK tool.Capture); // full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。 ToolOutputDir: cmdOutDir(runDir), - Compaction: compactionConfig(w.window), // long tool-heavy runs stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 + Compaction: compactionConfig(w.compactionWindow()), // long tool-heavy runs stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 } if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic) opts.Hooks = hooks @@ -371,6 +382,18 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. s := agentcore.NewSession(opts) defer s.Close() // release the session's background-task manager (temp dir + processes) + // Resume prior conversation if this intent was paused/blocked/exhausted and is + // being re-run. The transcript ID is deterministic per intent, so if a prior + // session exists the worker continues from where it left off instead of + // restarting from scratch. + if w.tx != nil { + _ = s.Resume(opts.SessionID) + if len(s.Messages()) > 0 { + seedUnlockFromHistory(s.Messages(), def.UnlockSkill) + input = "继续执行。" + } + } + // Budgets + settlement are owned by the SDK (MaxTurns/MaxDuration + Settlement): // on hit it runs a wrap-up turn and finishes with ReasonMaxTurns/ReasonTimeout. // We only add a HARD ctx backstop past the soft wall-clock budget, for the rare @@ -382,7 +405,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. runCtx := ctx if maxDur > 0 { var runCancel context.CancelFunc - runCtx, runCancel = context.WithTimeout(ctx, maxDur+settleHardGrace) + runCtx, runCancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout) defer runCancel() } _, reason, err := captureRunSession(runCtx, s, input, emitWrap) diff --git a/cmd/artex/main.go b/cmd/artex/main.go index 30440d4..cbc5b42 100644 --- a/cmd/artex/main.go +++ b/cmd/artex/main.go @@ -16,6 +16,7 @@ import ( "syscall" "time" + "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/config" "github.com/Autumn-27/artex/server" ) @@ -70,8 +71,10 @@ func main() { log.Printf("[config] 配置文件: %s (不存在 — 将仅尝试环境变量 ARTEX_PG_DSN)", cfgPath) } - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + sigCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() + ctx, shutdown := shutdownContext(sigCtx) + defer shutdown(agent.AbortShutdown) mgr, err := server.NewManager(*dataDir, *proxy) if err != nil { @@ -104,3 +107,18 @@ func main() { defer cancel() _ = httpSrv.Shutdown(shutdownCtx) } + +// shutdownContext deliberately does not derive from signalCtx. If it did, the +// parent's plain context.Canceled could win the race before AbortShutdown was +// attached to the child, losing the diagnostic cause in every running Agent. +func shutdownContext(signalCtx context.Context) (context.Context, context.CancelCauseFunc) { + ctx, shutdown := context.WithCancelCause(context.Background()) + go func() { + select { + case <-signalCtx.Done(): + shutdown(agent.AbortShutdown) + case <-ctx.Done(): + } + }() + return ctx, shutdown +} diff --git a/cmd/artex/main_test.go b/cmd/artex/main_test.go new file mode 100644 index 0000000..149c7c6 --- /dev/null +++ b/cmd/artex/main_test.go @@ -0,0 +1,19 @@ +package main + +import ( + "context" + "testing" + + "github.com/Autumn-27/artex/agent" +) + +func TestShutdownContextPreservesNamedCause(t *testing.T) { + signalCtx, signalCancel := context.WithCancel(context.Background()) + ctx, shutdown := shutdownContext(signalCtx) + defer shutdown(nil) + signalCancel() + <-ctx.Done() + if code, _, _, ok := agent.AbortReason(ctx); !ok || code != "shutdown" { + t.Fatalf("code=%q ok=%v, want shutdown", code, ok) + } +} diff --git a/db/db.go b/db/db.go index 8332b8e..de0a738 100644 --- a/db/db.go +++ b/db/db.go @@ -103,7 +103,7 @@ type builtinAgent struct { key, name, role, desc string vars []promptVar interactiveShell bool // 建行时的默认交互式 shell 开关；ON CONFLICT 不覆盖用户后续手动开关 - runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用表默认(600)，0=不限时 + runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用种子默认(1200)，0=不限时 } type promptVar struct{ name, desc, example, source string } @@ -145,7 +145,7 @@ func (d *DB) seedBuiltins() error { var agentID int64 err := d.QueryRow(` INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds) -VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 600)) +VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 1200)) ON CONFLICT (key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID) if err != nil { diff --git a/db/exploration.go b/db/exploration.go index fa54151..8747fa6 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -3,11 +3,17 @@ package db import ( "database/sql" "encoding/json" + "errors" "fmt" + "sort" "strings" "time" ) +// ErrIntentStateConflict marks a failed compare-and-set transition. Callers use +// it to distinguish an expected control race (HTTP 409) from a storage failure. +var ErrIntentStateConflict = errors.New("intent state changed concurrently") + // utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces // invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes. // Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's @@ -25,34 +31,40 @@ func utf8Clean(s string) string { // Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint. type Node struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - Payload json.RawMessage `json:"payload"` - Priority int `json:"priority"` - State string `json:"state"` - Origin string `json:"origin,omitempty"` - Owner string `json:"owner,omitempty"` - Anchors []int64 `json:"anchors,omitempty"` - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + Kind string `json:"kind"` + Payload json.RawMessage `json:"payload"` + Priority int `json:"priority"` + State string `json:"state"` + Origin string `json:"origin,omitempty"` + Owner string `json:"owner,omitempty"` + BlockedReason string `json:"blocked_reason,omitempty"` + Anchors []int64 `json:"anchors,omitempty"` + CreatedAt time.Time `json:"created_at"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // Activity is one worker execution step (= old task_activity). type Activity struct { - ID int64 `json:"id"` - NodeID *int64 `json:"node_id,omitempty"` - Worker string `json:"worker,omitempty"` - Kind string `json:"kind,omitempty"` - Tool string `json:"tool,omitempty"` - ToolUseID string `json:"tool_use_id,omitempty"` - IsError bool `json:"is_error"` - Summary string `json:"summary,omitempty"` - Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + NodeID *int64 `json:"node_id,omitempty"` + Worker string `json:"worker,omitempty"` + Kind string `json:"kind,omitempty"` + Tool string `json:"tool,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + IsError bool `json:"is_error"` + Summary string `json:"summary,omitempty"` + Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads + Metadata json.RawMessage `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` // Token usage — set only on the terminal kind='result' record; nil elsewhere. - InputTokens *int `json:"input_tokens,omitempty"` - OutputTokens *int `json:"output_tokens,omitempty"` - CacheReadTokens *int `json:"cache_read_tokens,omitempty"` - CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + InputTokens *int `json:"input_tokens,omitempty"` + OutputTokens *int `json:"output_tokens,omitempty"` + CacheReadTokens *int `json:"cache_read_tokens,omitempty"` + CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // TokenUsage is a per-worker token aggregate (TokenStatsByWorker). @@ -64,6 +76,17 @@ type TokenUsage struct { CacheWriteTokens int `json:"cache_write_tokens"` } +// SessionTokenUsage is the authoritative completed-run total for one task UI +// session. Worker sessions are keyed by intent rather than the reusable work#N +// executor name, so retries and reassignment remain attached to the same row. +type SessionTokenUsage struct { + Session string `json:"session"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheReadTokens int `json:"cache_read_tokens"` + CacheWriteTokens int `json:"cache_write_tokens"` +} + // DailyTokenBucket is one day's global token aggregate across all tasks. type DailyTokenBucket struct { Day string `json:"day"` // "YYYY-MM-DD" @@ -201,13 +224,10 @@ ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, re // SetNodeState updates any node's state (never deletes). func (s *ExplorationStore) SetNodeState(id int64, state string) error { - _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) + _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1, blocked_reason=NULL WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) return err } -<<<<<<< Updated upstream -// SetIntentState updates an intent node's state (done/blocked/exhausted/open). -======= // UpdateGoalPayload rewrites a goal node's payload text (and optional vulnclass); // scoped to kind='goal' so it can never mutate an intent/fact by id. Returns an // error if no such goal exists in this exploration. @@ -243,13 +263,12 @@ func (s *ExplorationStore) DeleteGoal(id int64) error { } // SetIntentState updates an intent node's state. ->>>>>>> Stashed changes // ResetRunningIntents returns any intent left in 'running' back to 'open' so it // is re-claimed. Called on startup: a 'running' intent with no live worker (a // backend restart or crashed worker goroutine left it stuck) would otherwise spin -// forever in the UI. Workers re-run an intent from scratch, so reopening is safe. +// forever in the UI. Workers resume from their transcript, so reopening is safe. func (s *ExplorationStore) ResetRunningIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) if err != nil { return 0, err @@ -259,12 +278,12 @@ WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) } // ReopenIntent flips ONE not-successfully-finished intent (blocked/exhausted/stopped) -// back to 'open' so a worker re-claims and re-runs it from scratch (graph writes it -// already produced stay). done/open/running are left untouched. Returns whether a row -// changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button. +// back to 'open' so a worker re-claims it (graph writes it already produced stay). +// The worker resumes from its prior LLM transcript instead of restarting from scratch. +// done/open/running are left untouched. Returns whether a row changed. Used by "重跑". func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) { res, err := s.db.Exec(`UPDATE exploration_nodes -SET state='open', completed_at=NULL +SET state='open', completed_at=NULL, blocked_reason=NULL WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state IN ('blocked','exhausted','stopped')`, id, s.expID) if err != nil { @@ -276,9 +295,9 @@ WHERE id=$1 AND exploration_id=$2 AND kind='intent' // ReopenBlockedIntents flips EVERY 'blocked' intent in this exploration back to 'open' // (batch rerun after e.g. an LLM/network outage that blocked many at once). Returns the -// number reopened. Workers re-run each from scratch; kept graph writes remain. +// number reopened. Workers resume from their transcripts; kept graph writes remain. func (s *ExplorationStore) ReopenBlockedIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID) if err != nil { return 0, err @@ -291,14 +310,11 @@ func (s *ExplorationStore) SetIntentState(id int64, state string) error { // terminal states stamp completed_at; reopening (back to open/running) clears it. terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped" _, err := s.db.Exec(`UPDATE exploration_nodes -SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END +SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $4 THEN now() ELSE NULL END WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal) return err } -<<<<<<< Updated upstream -const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at` -======= // CompareAndSetIntentState transitions one local intent only when its current // state still matches expected. It is the state boundary used by pause/resume and // worker settlement, so a stale API read cannot overwrite a concurrent finish. @@ -593,12 +609,11 @@ func (u *TokenUsage) add(other TokenUsage) { } const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at` ->>>>>>> Stashed changes func scanNode(sc interface{ Scan(...any) error }) (*Node, error) { var n Node var payload []byte - if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil { + if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.BlockedReason, &n.CreatedAt); err != nil { return nil, err } n.Payload = json.RawMessage(payload) @@ -790,6 +805,31 @@ WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields return out, rows.Err() } +// FindingIntentsTerminal is the inherited-history variant: it returns finding +// lineage only when the producing intent has reached an immutable terminal +// state. This keeps a source task's live worker topology private. +func (s *ExplorationStore) FindingIntentsTerminal() (map[int64]int64, error) { + rows, err := s.db.Query(`SELECT e.dst_id, e.src_id + FROM exploration_edges e + JOIN exploration_nodes finding ON finding.id=e.dst_id AND finding.exploration_id=e.exploration_id + JOIN exploration_nodes intent ON intent.id=e.src_id AND intent.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.rel=$2 AND finding.kind='finding' + AND intent.kind='intent' AND intent.state IN ('done','blocked','exhausted','stopped')`, s.expID, RelYields) + if err != nil { + return nil, err + } + defer rows.Close() + out := map[int64]int64{} + for rows.Next() { + var findingID, intentID int64 + if err := rows.Scan(&findingID, &intentID); err != nil { + return nil, err + } + out[findingID] = intentID + } + return out, rows.Err() +} + // Frontier returns open intents ordered by priority desc then id (FIFO tiebreak). func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) { if limit <= 0 { @@ -851,12 +891,16 @@ func (s *ExplorationStore) Stats() (map[string]int, error) { // AppendActivity records one worker step and returns its id. func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) { + metadata := a.Metadata + if len(metadata) == 0 { + metadata = json.RawMessage(`{}`) + } var id int64 err := s.db.QueryRow(` -INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13) +INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14) RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail), - a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) + metadata, a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) return id, err } @@ -986,6 +1030,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) { return out, rows.Err() } +// TokenStatsBySession aggregates every persisted completed run, independent of +// activity-history pagination. Main/planner use fixed keys; Worker executions use +// their intent node id, even when a different work#N process handled a retry. +func (s *ExplorationStore) TokenStatsBySession() ([]SessionTokenUsage, error) { + rows, err := s.db.Query(`SELECT + CASE + WHEN COALESCE(a.worker,'')='mainagent' THEN 'main' + WHEN COALESCE(a.worker,'')='planner' THEN 'plan' + WHEN n.id IS NOT NULL THEN 'intent:' || n.id::text + ELSE '' + END AS session_key, + COALESCE(SUM(a.input_tokens),0), COALESCE(SUM(a.output_tokens),0), + COALESCE(SUM(a.cache_read_tokens),0), COALESCE(SUM(a.cache_write_tokens),0) + FROM activity a + LEFT JOIN exploration_nodes n + ON n.id=a.node_id AND n.exploration_id=a.exploration_id AND n.kind='intent' + WHERE a.exploration_id=$1 AND a.kind='result' + AND (COALESCE(a.worker,'') IN ('mainagent','planner') OR n.id IS NOT NULL) + GROUP BY session_key + ORDER BY session_key`, s.expID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []SessionTokenUsage{} + for rows.Next() { + var u SessionTokenUsage + if err := rows.Scan(&u.Session, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil { + return nil, err + } + out = append(out, u) + } + return out, rows.Err() +} + // ActivityList returns steps after sinceID (exclusive), optionally filtered by node. // Returns items and the new cursor (max id seen). func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) { @@ -994,7 +1073,7 @@ func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) } var rows *sql.Rows var err error - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` if nodeID != nil { rows, err = s.db.Query(`SELECT `+cols+` FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit) @@ -1008,6 +1087,43 @@ FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID, defer rows.Close() out := []Activity{} cursor := sinceID + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, sinceID, err + } + if a.ID > cursor { + cursor = a.ID + } + out = append(out, a) + } + return out, cursor, rows.Err() +} + +// ActivityListForTerminalIntent is the inherited, incremental-read boundary. +// The intent state check and activity read happen in one statement so a source +// intent reopened between API-level checks cannot expose its new in-flight trace. +// Model reasoning and accounting rows are never inherited. +func (s *ExplorationStore) ActivityListForTerminalIntent(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + if limit <= 0 { + limit = 300 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND a.id>$3 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $4`, s.expID, nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + defer rows.Close() + out := []Activity{} + cursor := sinceID for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -1054,7 +1170,7 @@ func (s *ExplorationStore) ActivityPage(f ActivitySessionFilter, before int64, l if limit <= 0 { limit = 200 } - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` args := []any{s.expID} cond, cargs := f.cond(len(args) + 1) args = append(args, cargs...) @@ -1071,6 +1187,51 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) } defer rows.Close() desc := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, false, err + } + desc = append(desc, a) + } + if err := rows.Err(); err != nil { + return nil, false, err + } + hasMore := len(desc) > limit + if hasMore { + desc = desc[:limit] + } + // reverse the newest-first window into ascending id order for display. + out := make([]Activity, len(desc)) + for i, a := range desc { + out[len(desc)-1-i] = a + } + return out, hasMore, nil +} + +// ActivityPageForTerminalIntent is the inherited session-history boundary. It +// combines ownership, terminal-state and row-kind checks in one query, closing +// the race where a previously completed source intent is reopened before paging. +func (s *ExplorationStore) ActivityPageForTerminalIntent(nodeID, before int64, limit int) ([]Activity, bool, error) { + if limit <= 0 { + limit = 200 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND ($3 <= 0 OR a.id < $3) + ORDER BY a.id DESC LIMIT $4`, s.expID, nodeID, before, limit+1) + if err != nil { + return nil, false, err + } + defer rows.Close() + desc := []Activity{} for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -1086,7 +1247,6 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) if hasMore { desc = desc[:limit] } - // reverse the newest-first window into ascending id order for display. out := make([]Activity, len(desc)) for i, a := range desc { out[len(desc)-1-i] = a @@ -1152,6 +1312,26 @@ ORDER BY id LIMIT $3`, s.expID, nodeID, limit) return scanTrace(rows) } +// ActivityTraceForTerminalIntent returns one inherited work trace only while its +// source intent is still terminal at query time. +func (s *ExplorationStore) ActivityTraceForTerminalIntent(nodeID int64, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 500 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $3`, s.expID, nodeID, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // ActivityTraceSearch finds steps whose summary OR detail matches q // (case-insensitive), returning summaries only. nodeID != nil scopes to one work; // nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no @@ -1179,13 +1359,59 @@ AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like, return scanTrace(rows) } +// ActivityTraceSearchForTerminalIntent is the scoped inherited search variant; +// terminal eligibility is evaluated by the same statement that reads the rows. +func (s *ExplorationStore) ActivityTraceSearchForTerminalIntent(nodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $3 OR a.detail ILIKE $3) + ORDER BY a.id LIMIT $4`, s.expID, nodeID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + +// ActivityTraceSearchTerminalIntents searches inherited worker history without +// exposing planner/main rows or work that is still open/running in the source. +func (s *ExplorationStore) ActivityTraceSearchTerminalIntents(q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $2 OR a.detail ILIKE $2) + ORDER BY a.id LIMIT $3`, s.expID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // AssetRef is a compact exploration node (intent/fact/finding) that references an // asset via an anchor — for the coverage graph's node drawer. type AssetRef struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - State string `json:"state"` - Summary string `json:"summary"` + ID int64 `json:"id"` + Kind string `json:"kind"` + State string `json:"state"` + Summary string `json:"summary"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // AssetRefs returns the intents / facts / findings in THIS exploration anchored to @@ -1278,6 +1504,42 @@ FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings. return out, rows.Err() } +// ActivityByIDsForTerminalIntents is the inherited-detail read boundary. It +// deliberately excludes node-less planner/main rows, non-intent activities, live +// source work, and thinking. Local callers that need legacy behavior use +// ActivityByIDs instead. +func (s *ExplorationStore) ActivityByIDsForTerminalIntents(ids []int64) ([]Activity, error) { + if len(ids) == 0 { + return nil, nil + } + ph := make([]string, len(ids)) + args := make([]any, len(ids)+1) + args[0] = s.expID + for i, id := range ids { + ph[i] = fmt.Sprintf("$%d", i+2) + args[i+1] = id + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.detail,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.kind NOT IN ('thinking','usage') AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.id IN (`+strings.Join(ph, ",")+`) ORDER BY a.id`, args...) + if err != nil { + return nil, err + } + defer rows.Close() + out := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil { + return nil, err + } + out = append(out, a) + } + return out, rows.Err() +} + // AddStandaloneFinding writes a finding to the standalone findings table, which // persists across task deletion (task_id / node_id become NULL when the task or // exploration node is deleted). taskID and nodeID may be 0 (stored as NULL). diff --git a/db/exploration_sources.go b/db/exploration_sources.go new file mode 100644 index 0000000..3ccbb96 --- /dev/null +++ b/db/exploration_sources.go @@ -0,0 +1,350 @@ +package db + +import ( + "database/sql" + "sort" +) + +// DirectSourceStore binds one directly related task to its exploration store. +// It is intentionally a read-side helper: callers keep using the receiver store +// for every graph mutation, frontier lookup, and intent claim. +type DirectSourceStore struct { + Task TaskSource + Store *ExplorationStore +} + +// DirectSourceStores resolves the live, direct task relations for this +// exploration. It deliberately queries on every call: relations disappear when +// a source task is deleted, and inherited context must not retain stale rows. +// Sources of a source are never expanded. +func (s *ExplorationStore) DirectSourceStores() ([]DirectSourceStore, error) { + rows, err := s.db.Query(` +SELECT source.id, source.exploration_id, source.description, source.goal, source.status +FROM tasks owner +JOIN task_relations relation ON relation.task_id=owner.id +JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL +WHERE owner.exploration_id=$1 AND owner.deleted_at IS NULL +ORDER BY relation.created_at, source.id`, s.expID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []DirectSourceStore{} + for rows.Next() { + var source TaskSource + if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil { + return nil, err + } + out = append(out, DirectSourceStore{Task: source, Store: s.db.Exploration(source.ExplorationID)}) + } + return out, rows.Err() +} + +// TaskID returns the live task bound to this exploration. Explorations created +// directly in tests or maintenance code have no task and return zero. +func (s *ExplorationStore) TaskID() (int64, error) { + var id int64 + err := s.db.QueryRow(`SELECT id FROM tasks WHERE exploration_id=$1 AND deleted_at IS NULL`, s.expID).Scan(&id) + if err != nil { + if err == sql.ErrNoRows { + return 0, nil + } + return 0, err + } + return id, nil +} + +func markInheritedNode(n *Node, taskID int64) *Node { + if n == nil { + return nil + } + n.SourceTaskID = taskID + n.Inherited = true + return n +} + +func markInheritedActivities(in []Activity, taskID int64) []Activity { + for i := range in { + in[i].SourceTaskID = taskID + in[i].Inherited = true + } + return in +} + +func inheritedIntentTerminal(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false + } +} + +// ListByKindWithSources returns local nodes followed by nodes from each direct +// source in relation order. The limit remains per exploration, matching the +// existing ListByKind contract while ensuring one large task cannot hide all +// inherited context from another source. +func (s *ExplorationStore) ListByKindWithSources(kind string, limit int) ([]*Node, error) { + out, err := s.ListByKind(kind, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + nodes, err := source.Store.ListByKind(kind, limit) + if err != nil { + return nil, err + } + for _, node := range nodes { + if kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + out = append(out, markInheritedNode(node, source.Task.TaskID)) + } + } + return out, nil +} + +// GetNodeWithSources reads a node only when it belongs to this exploration or +// one of its direct sources. Inherited nodes are tagged so tool callers can keep +// them read-only and show their provenance. +func (s *ExplorationStore) GetNodeWithSources(id int64) (*Node, error) { + node, err := s.GetNode(id) + if err != nil || node != nil { + return node, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, err = source.Store.GetNode(id) + if err != nil { + return nil, err + } + if node != nil { + if node.Kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + return markInheritedNode(node, source.Task.TaskID), nil + } + } + return nil, nil +} + +// FindingIntentsWithSources combines finding lineage for the current +// exploration and each direct source. Node ids are globally unique. +func (s *ExplorationStore) FindingIntentsWithSources() (map[int64]int64, error) { + out, err := s.FindingIntents() + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + items, err := source.Store.FindingIntentsTerminal() + if err != nil { + return nil, err + } + for findingID, intentID := range items { + out[findingID] = intentID + } + } + return out, nil +} + +// ActivityTraceWithSources returns a work trace when its intent belongs to the +// current exploration or a direct source. It never searches indirect sources. +func (s *ExplorationStore) ActivityTraceWithSources(nodeID int64, limit int) ([]Activity, error) { + acts, err := s.ActivityTrace(nodeID, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceForTerminalIntent(nodeID, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityListWithSources is the source-aware equivalent used by +// get_worker_output. Node ids are global, so the first owning exploration is +// unambiguous even when the work has not emitted any activity yet. +func (s *ExplorationStore) ActivityListWithSources(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + node, err := s.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node != nil { + return s.ActivityList(&nodeID, sinceID, limit) + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, sinceID, err + } + for _, source := range sources { + node, err = source.Store.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, cursor, err := source.Store.ActivityListForTerminalIntent(nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + return markInheritedActivities(acts, source.Task.TaskID), cursor, nil + } + return []Activity{}, sinceID, nil +} + +// ActivityDetailWithSources keeps the legacy local-task lookup (including local +// thinking rows), while inherited details are restricted to terminal worker +// intents. Source planner/main rows have no node id and must never become part of +// inherited context. +func (s *ExplorationStore) ActivityDetailWithSources(id int64) (string, error) { + detail, err := s.ActivityDetail(id) + if err != nil || detail != "" { + return detail, err + } + acts, err := s.ActivityByIDsWithSources([]int64{id}) + if err != nil || len(acts) == 0 { + return "", err + } + return acts[0].Detail, nil +} + +// ActivityTraceSearchWithSources performs the scoped worker-trace search used +// by get_worker_trace. A node id resolves to at most one exploration because +// exploration node ids are global. +func (s *ExplorationStore) ActivityTraceSearchWithSources(nodeID int64, q string, limit int) ([]Activity, error) { + acts, err := s.ActivityTraceSearch(&nodeID, q, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceSearchForTerminalIntent(nodeID, q, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityTraceSearchAllWithSources searches local worker traces plus every +// direct source. The local owner is excluded only from the current exploration; +// inherited traces are immutable historical context. +func (s *ExplorationStore) ActivityTraceSearchAllWithSources(excludeNodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + out, err := s.ActivityTraceSearchExcluding(excludeNodeID, q, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityTraceSearchTerminalIntents(q, limit) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +// ActivityByIDsWithSources loads local step details with the legacy behavior. For +// direct sources it only returns rows attached to terminal intents, preventing an +// arbitrary global activity id from exposing source planner/main transcripts. +func (s *ExplorationStore) ActivityByIDsWithSources(ids []int64) ([]Activity, error) { + out, err := s.ActivityByIDs(ids) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityByIDsForTerminalIntents(ids) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + return out, nil +} + +// AssetRefsWithSources returns anchored nodes from this exploration and each +// direct source. Inherited entries retain their owning task id so API/UI callers +// can present them as immutable context. +func (s *ExplorationStore) AssetRefsWithSources(assetID int64) ([]AssetRef, error) { + out, err := s.AssetRefs(assetID) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + refs, err := source.Store.AssetRefs(assetID) + if err != nil { + return nil, err + } + for i := range refs { + if refs[i].Kind == KindIntent && !inheritedIntentTerminal(refs[i].State) { + continue + } + refs[i].SourceTaskID = source.Task.TaskID + refs[i].Inherited = true + out = append(out, refs[i]) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].ID > out[j].ID }) + return out, nil +} diff --git a/db/exploration_sources_test.go b/db/exploration_sources_test.go new file mode 100644 index 0000000..07ceb29 --- /dev/null +++ b/db/exploration_sources_test.go @@ -0,0 +1,431 @@ +package db + +import ( + "fmt" + "testing" +) + +func TestInheritedActivityReadsRequireTerminalIntent(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("source activity boundary", "source activity boundary") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) }) + store := d.Exploration(expID) + intentID, err := store.AddIntent(map[string]any{"summary": "source work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + appendStep := func(kind, summary string) int64 { + t.Helper() + id, appendErr := store.AppendActivity(Activity{ + NodeID: &intentID, Worker: "worker", Kind: kind, Summary: summary, Detail: summary + " detail", + }) + if appendErr != nil { + t.Fatal(appendErr) + } + return id + } + textID := appendStep("text", "shared text") + thinkingID := appendStep("thinking", "private reasoning") + usageID := appendStep("usage", "private accounting") + resultID := appendStep("result", "shared result") + if err := store.SetIntentState(intentID, "done"); err != nil { + t.Fatal(err) + } + + page, more, err := store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || more || len(page) != 2 || page[0].ID != textID || page[1].ID != resultID { + t.Fatalf("terminal page boundary: page=%+v more=%v err=%v", page, more, err) + } + list, cursor, err := store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 2 || cursor != resultID { + t.Fatalf("terminal list boundary: list=%+v cursor=%d err=%v", list, cursor, err) + } + trace, err := store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 2 { + t.Fatalf("terminal trace boundary: trace=%+v err=%v", trace, err) + } + hits, err := store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 2 { + t.Fatalf("terminal search boundary: hits=%+v err=%v", hits, err) + } + details, err := store.ActivityByIDsForTerminalIntents([]int64{textID, thinkingID, usageID, resultID}) + if err != nil || len(details) != 2 || details[0].ID != textID || details[1].ID != resultID { + t.Fatalf("terminal detail boundary: details=%+v err=%v", details, err) + } + + // Reopening a source intent must close every inherited activity read even if + // callers still hold a stale terminal-state snapshot. + if err := store.SetIntentState(intentID, "running"); err != nil { + t.Fatal(err) + } + page, _, err = store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || len(page) != 0 { + t.Fatalf("reopened intent leaked through page: page=%+v err=%v", page, err) + } + list, _, err = store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 0 { + t.Fatalf("reopened intent leaked through list: list=%+v err=%v", list, err) + } + trace, err = store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 0 { + t.Fatalf("reopened intent leaked through trace: trace=%+v err=%v", trace, err) + } + hits, err = store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 0 { + t.Fatalf("reopened intent leaked through search: hits=%+v err=%v", hits, err) + } + details, err = store.ActivityByIDsForTerminalIntents([]int64{textID, resultID}) + if err != nil || len(details) != 0 { + t.Fatalf("reopened intent leaked through detail: details=%+v err=%v", details, err) + } +} + +func TestExplorationDirectSourceReadView(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand source", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("direct source", "source goal", TaskCreateOptions{ + SourceTaskIDs: []int64{grand.ID}, + }) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current", "current goal", TaskCreateOptions{ + SourceTaskIDs: []int64{source.ID}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + }) + + grandStore := d.Exploration(grand.ExplorationID) + sourceStore := d.Exploration(source.ExplorationID) + currentStore := d.Exploration(current.ExplorationID) + + grandFact, err := grandStore.AddNode(KindFact, map[string]any{"summary": "grand-only fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFact, err := sourceStore.AddNode(KindFact, map[string]any{"summary": "direct source fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFinding, err := sourceStore.AddNode(KindFinding, map[string]any{"summary": "direct source finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work"}, 9, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.Link(sourceIntent, RelYields, sourceFinding); err != nil { + t.Fatal(err) + } + currentIntent, err := currentStore.AddIntent(map[string]any{"summary": "current work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + currentFact, err := currentStore.AddNode(KindFact, map[string]any{"summary": "current fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + + sources, err := currentStore.DirectSourceStores() + if err != nil { + t.Fatal(err) + } + if len(sources) != 1 || sources[0].Task.TaskID != source.ID { + t.Fatalf("want only direct source %d, got %+v", source.ID, sources) + } + + frontier, err := currentStore.Frontier(20) + if err != nil { + t.Fatal(err) + } + if len(frontier) != 1 || frontier[0].ID != currentIntent { + t.Fatalf("source intent leaked into local frontier: %+v", frontier) + } + if claimed, err := currentStore.ClaimIntent(sourceIntent, "wrong-task-worker"); err != nil || claimed { + t.Fatalf("source intent must not be claimable: claimed=%v err=%v", claimed, err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got != nil { + t.Fatalf("live source intent must not be readable as inherited history: node=%+v err=%v", got, err) + } + if got, err := currentStore.ListByKindWithSources(KindIntent, 100); err != nil || len(got) != 1 || got[0].ID != currentIntent { + t.Fatalf("live source intent leaked through inherited list: nodes=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != 0 { + t.Fatalf("live source intent leaked through finding lineage: lineage=%+v err=%v", got, err) + } + + facts, err := currentStore.ListByKindWithSources(KindFact, 100) + if err != nil { + t.Fatal(err) + } + seenCurrent, seenSource, seenGrand := false, false, false + for _, fact := range facts { + switch fact.ID { + case currentFact: + seenCurrent = !fact.Inherited && fact.SourceTaskID == 0 + case sourceFact: + seenSource = fact.Inherited && fact.SourceTaskID == source.ID + case grandFact: + seenGrand = true + } + } + if !seenCurrent || !seenSource || seenGrand { + t.Fatalf("unexpected direct-source fact view: current=%v source=%v grand=%v", seenCurrent, seenSource, seenGrand) + } + + gotSource, err := currentStore.GetNodeWithSources(sourceFact) + if err != nil || gotSource == nil || !gotSource.Inherited || gotSource.SourceTaskID != source.ID { + t.Fatalf("source node lookup: node=%+v err=%v", gotSource, err) + } + gotGrand, err := currentStore.GetNodeWithSources(grandFact) + if err != nil || gotGrand != nil { + t.Fatalf("indirect source must be invisible: node=%+v err=%v", gotGrand, err) + } + + // Existing write methods remain bound to currentStore.expID. Passing an + // inherited id is a no-op and cannot mutate the source blackboard. + if err := currentStore.SetNodeState(sourceFact, "dismissed"); err != nil { + t.Fatal(err) + } + unchanged, err := sourceStore.GetNode(sourceFact) + if err != nil || unchanged == nil || unchanged.State != "confirmed" { + t.Fatalf("inherited fact was mutated: node=%+v err=%v", unchanged, err) + } + + if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil { + t.Fatal(err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got == nil || !got.Inherited { + t.Fatalf("terminal source intent should become readable history: node=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != sourceIntent { + t.Fatalf("terminal source finding lineage missing: lineage=%+v err=%v", got, err) + } + sourcePlannerStep, err := sourceStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "private source planner", Detail: "private source planner transcript", + }) + if err != nil { + t.Fatal(err) + } + openSourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work still running"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.SetIntentState(openSourceIntent, "running"); err != nil { + t.Fatal(err) + } + openSourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &openSourceIntent, Worker: "source-live-worker", Kind: "text", + Summary: "private live source work", Detail: "private live source transcript", + }) + if err != nil { + t.Fatal(err) + } + sourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &sourceIntent, Worker: "source-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "source-secret response", Detail: "source-secret full detail", + }) + if err != nil { + t.Fatal(err) + } + currentPlannerStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "current planner", Detail: "current planner detail", + }) + if err != nil { + t.Fatal(err) + } + currentThinkingStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "thinking", Summary: "current thinking", Detail: "current thinking detail", + }) + if err != nil { + t.Fatal(err) + } + grandStep, err := grandStore.AppendActivity(Activity{ + NodeID: &grandFact, Worker: "grand-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "grand-only response", Detail: "grand-only detail", + }) + if err != nil { + t.Fatal(err) + } + + trace, err := currentStore.ActivityTraceWithSources(sourceIntent, 20) + if err != nil || len(trace) != 1 || trace[0].ID != sourceStep || !trace[0].Inherited || trace[0].SourceTaskID != source.ID { + t.Fatalf("source trace lookup: trace=%+v err=%v", trace, err) + } + hits, err := currentStore.ActivityTraceSearchAllWithSources(0, "source-secret", 20) + if err != nil || len(hits) != 1 || hits[0].ID != sourceStep || !hits[0].Inherited { + t.Fatalf("source trace search: hits=%+v err=%v", hits, err) + } + grandHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "grand-only", 20) + if err != nil || len(grandHits) != 0 { + t.Fatalf("indirect source trace leaked: hits=%+v err=%v", grandHits, err) + } + liveHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "private live source", 20) + if err != nil || len(liveHits) != 0 { + t.Fatalf("live source trace leaked: hits=%+v err=%v", liveHits, err) + } + if liveTrace, err := currentStore.ActivityTraceWithSources(openSourceIntent, 20); err != nil || len(liveTrace) != 0 { + t.Fatalf("live inherited intent trace leaked: trace=%+v err=%v", liveTrace, err) + } + details, err := currentStore.ActivityByIDsWithSources([]int64{ + sourcePlannerStep, openSourceStep, sourceStep, currentPlannerStep, grandStep, + }) + if err != nil || len(details) != 2 || details[0].ID != sourceStep || !details[0].Inherited || details[1].ID != currentPlannerStep || details[1].Inherited { + t.Fatalf("source trace detail allow-list: details=%+v err=%v", details, err) + } + if detail, err := currentStore.ActivityDetailWithSources(sourcePlannerStep); err != nil || detail != "" { + t.Fatalf("source planner transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(openSourceStep); err != nil || detail != "" { + t.Fatalf("live source intent transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentPlannerStep); err != nil || detail != "current planner detail" { + t.Fatalf("local activity detail behavior changed: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentThinkingStep); err != nil || detail != "current thinking detail" { + t.Fatalf("local thinking detail behavior changed: detail=%q err=%v", detail, err) + } +} + +func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand assets", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("source assets", "source goal", TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}}) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current assets", "current goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + sourceTestedDomain := fmt.Sprintf("context-tested-%d.invalid", source.ID) + sourceUntestedDomain := fmt.Sprintf("context-untested-%d.invalid", source.ID) + sourceAnchoredOnlyDomain := fmt.Sprintf("context-anchor-only-%d.invalid", source.ID) + grandDomain := fmt.Sprintf("context-grand-%d.invalid", grand.ID) + sourceTestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceTestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceUntestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceUntestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceAnchoredOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceAnchoredOnlyDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + grandID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: grandDomain, TaskID: grand.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + _, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID}) + }) + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test", "agent"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test", "agent"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test", "agent"); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "tested source asset"}, 0, "confirmed", "worker", []int64{sourceTestedID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "anchored-only source asset"}, 0, "confirmed", "worker", []int64{sourceAnchoredOnlyID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(grand.ExplorationID).AddNode(KindFact, map[string]any{"summary": "indirect tested asset"}, 0, "confirmed", "worker", []int64{grandID}); err != nil { + t.Fatal(err) + } + + scopes, err := assets.ListTaskScopeWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if len(scopes) != 2 || scopes[0].TaskID != source.ID || scopes[1].TaskID != source.ID { + t.Fatalf("direct source scopes only: %+v", scopes) + } + cov, err := assets.TaskCoverageWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if cov.Denominator != 3 || cov.Tested != 2 { + t.Fatalf("combined coverage: %+v", cov) + } + untested, total, err := assets.ListUntestedAssetsWithSources(current.ID, "", 10, 0) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(untested) != 1 || untested[0].ID != sourceUntestedID { + t.Fatalf("combined untested assets: total=%d assets=%+v", total, untested) + } + hosts, err := assets.HostsByTaskWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + hostSet := map[string]bool{} + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[sourceTestedDomain] || !hostSet[sourceUntestedDomain] || !hostSet[sourceAnchoredOnlyDomain] || hostSet[grandDomain] { + t.Fatalf("direct source hosts only: %v", hosts) + } + refs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(sourceAnchoredOnlyID) + if err != nil || len(refs) != 1 || !refs[0].Inherited || refs[0].SourceTaskID != source.ID || refs[0].Kind != KindFact { + t.Fatalf("source asset refs provenance: refs=%+v err=%v", refs, err) + } + grandRefs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(grandID) + if err != nil || len(grandRefs) != 0 { + t.Fatalf("indirect source asset refs leaked: refs=%+v err=%v", grandRefs, err) + } + graph, err := assets.BuildCoverageGraph(current.ID, current.ExplorationID) + if err != nil { + t.Fatal(err) + } + seenAnchoredOnly := false + for _, node := range graph.Nodes { + if node.AssetID == sourceAnchoredOnlyID { + seenAnchoredOnly = node.Tested && node.InScope + } + } + if !seenAnchoredOnly { + t.Fatalf("anchored-only source asset missing from coverage graph: %+v", graph.Nodes) + } +} diff --git a/db/exploration_test.go b/db/exploration_test.go index 1c31d9f..97d9f62 100644 --- a/db/exploration_test.go +++ b/db/exploration_test.go @@ -1,6 +1,9 @@ package db -import "testing" +import ( + "fmt" + "testing" +) func TestExplorationFlow(t *testing.T) { d, err := Open(testDSN(t)) @@ -112,3 +115,134 @@ func TestExplorationFlow(t *testing.T) { t.Fatalf("stats: %+v", st) } } + +func TestIntentPauseResumeAndCancelCleanup(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("test", "worker control cleanup") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + es := d.Exploration(expID) + + assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("cancel-intent-%d.invalid", expID), + }) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, assetID) + + otherIntent, err := es.AddIntent(map[string]any{"summary": "keep intent"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + intentID, err := es.AddIntent(map[string]any{"summary": "cancel intent"}, 10, nil, "planner") + if err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-control-test"); err != nil || !claimed { + t.Fatalf("initial claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + frontier, err := es.Frontier(10) + if err != nil { + t.Fatal(err) + } + for _, node := range frontier { + if node.ID == intentID { + t.Fatalf("paused intent %d must not enter frontier", intentID) + } + } + if claimed, err := es.ClaimIntent(intentID, "worker-while-paused"); err != nil || claimed { + t.Fatalf("paused intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "open"); err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-after-resume"); err != nil || !claimed { + t.Fatalf("resumed intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + directFact, err := es.AddNode("fact", map[string]any{"summary": "remove fact"}, 0, "confirmed", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + directFinding, err := es.AddNode("finding", map[string]any{"summary": "remove finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + keptFact, err := es.AddNode("fact", map[string]any{"summary": "keep fact"}, 0, "confirmed", "other-worker", nil) + if err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFact); err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFinding); err != nil { + t.Fatal(err) + } + findingRowID, err := es.AddStandaloneFinding(0, directFinding, "test", "cancelled finding", SeverityHigh, "summary", "evidence", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + activityID, err := es.AppendActivity(Activity{NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "remove activity"}) + if err != nil { + t.Fatal(err) + } + keptActivityID, err := es.AppendActivity(Activity{NodeID: &otherIntent, Worker: "other-worker", Kind: "result", Summary: "keep activity"}) + if err != nil { + t.Fatal(err) + } + + cleanup, err := es.CancelIntent(intentID) + if err != nil { + t.Fatal(err) + } + if cleanup.Intents != 1 || cleanup.Facts != 1 || cleanup.Findings != 1 || cleanup.Activities != 1 { + t.Fatalf("unexpected cleanup counts: %+v", cleanup) + } + for _, nodeID := range []int64{intentID, directFact, directFinding} { + node, err := es.GetNode(nodeID) + if err != nil { + t.Fatal(err) + } + if node != nil { + t.Fatalf("node %d survived intent cancellation", nodeID) + } + } + for _, nodeID := range []int64{otherIntent, keptFact} { + node, err := es.GetNode(nodeID) + if err != nil || node == nil { + t.Fatalf("unrelated node %d removed: node=%v err=%v", nodeID, node, err) + } + } + + assertCount := func(query string, want int, args ...any) { + t.Helper() + var got int + if err := d.QueryRow(query, args...).Scan(&got); err != nil { + t.Fatal(err) + } + if got != want { + t.Fatalf("query count=%d, want %d: %s", got, want, query) + } + } + assertCount(`SELECT COUNT(*) FROM findings WHERE id=$1`, 0, findingRowID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 0, activityID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 1, keptActivityID) + assertCount(`SELECT COUNT(*) FROM exploration_edges WHERE exploration_id=$1 AND (src_id=$2 OR dst_id=$2)`, 0, expID, intentID) + assertCount(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1`, 0, directFact) + assertCount(`SELECT COUNT(*) FROM assets WHERE id=$1`, 1, assetID) +} diff --git a/db/exploration_tokens_test.go b/db/exploration_tokens_test.go new file mode 100644 index 0000000..a0bbf16 --- /dev/null +++ b/db/exploration_tokens_test.go @@ -0,0 +1,130 @@ +package db + +import ( + "fmt" + "testing" +) + +func TestTokenStatsBySessionUsesCompleteIntentHistory(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("token sessions", "aggregate complete session history") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + store := d.Exploration(expID) + + intentID, err := store.AddIntent(map[string]any{"summary": "worker session"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + factID, err := store.AddNode(KindFact, map[string]any{"summary": "not a worker"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + + otherExpID, err := d.CreateExploration("other token sessions", "foreign intent must not match") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, otherExpID) + foreignIntentID, err := d.Exploration(otherExpID).AddIntent(map[string]any{"summary": "foreign worker"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + + appendResult := func(nodeID *int64, worker string, input, output, read, write int) { + t.Helper() + if _, appendErr := store.AppendActivity(Activity{ + NodeID: nodeID, + Worker: worker, + Kind: "result", + InputTokens: &input, + OutputTokens: &output, + CacheReadTokens: &read, + CacheWriteTokens: &write, + }); appendErr != nil { + t.Fatal(appendErr) + } + } + + appendResult(nil, "mainagent", 11, 12, 13, 14) + appendResult(nil, "planner", 21, 22, 23, 24) + // More than ActivityPage's default limit proves this aggregate reads the + // persisted history directly instead of summing the currently loaded page. + const completedRuns = 205 + for i := 0; i < completedRuns; i++ { + appendResult(&intentID, fmt.Sprintf("work#%d", i%3+1), 1, 2, 3, 4) + } + + // All three rows remain part of the legacy worker/whole-task totals, but none + // represents a local Worker session and therefore none may create intent:*. + appendResult(&factID, "work#fact", 31, 32, 33, 34) + appendResult(nil, "work#missing-node", 41, 42, 43, 44) + appendResult(&foreignIntentID, "work#foreign", 51, 52, 53, 54) + ignoredInput := 1000 + if _, err := store.AppendActivity(Activity{ + NodeID: &intentID, + Worker: "work#1", + Kind: "usage", + InputTokens: &ignoredInput, + }); err != nil { + t.Fatal(err) + } + + sessions, err := store.TokenStatsBySession() + if err != nil { + t.Fatal(err) + } + got := make(map[string]SessionTokenUsage, len(sessions)) + for _, usage := range sessions { + got[usage.Session] = usage + } + if len(got) != 3 { + t.Fatalf("sessions = %+v, want only main, plan, and the local intent", sessions) + } + assertSessionTokens(t, got["main"], "main", 11, 12, 13, 14) + assertSessionTokens(t, got["plan"], "plan", 21, 22, 23, 24) + assertSessionTokens(t, got[fmt.Sprintf("intent:%d", intentID)], fmt.Sprintf("intent:%d", intentID), + completedRuns, completedRuns*2, completedRuns*3, completedRuns*4) + + workers, err := store.TokenStatsByWorker() + if err != nil { + t.Fatal(err) + } + if len(workers) != 8 { // main, planner, 3 executors, and the 3 deliberately invalid session rows + t.Fatalf("legacy workers changed: got %d entries: %+v", len(workers), workers) + } + total, err := store.TokenTotal() + if err != nil { + t.Fatal(err) + } + assertTokenUsage(t, total, + 11+21+completedRuns+31+41+51, + 12+22+completedRuns*2+32+42+52, + 13+23+completedRuns*3+33+43+53, + 14+24+completedRuns*4+34+44+54) +} + +func assertSessionTokens(t *testing.T, got SessionTokenUsage, session string, input, output, read, write int) { + t.Helper() + if got.Session != session || got.InputTokens != input || got.OutputTokens != output || + got.CacheReadTokens != read || got.CacheWriteTokens != write { + t.Fatalf("session %q = %+v, want input=%d output=%d cache-read=%d cache-write=%d", + session, got, input, output, read, write) + } +} + +func assertTokenUsage(t *testing.T, got TokenUsage, input, output, read, write int) { + t.Helper() + if got.InputTokens != input || got.OutputTokens != output || + got.CacheReadTokens != read || got.CacheWriteTokens != write { + t.Fatalf("token total = %+v, want input=%d output=%d cache-read=%d cache-write=%d", + got, input, output, read, write) + } +} diff --git a/db/intent_admission.go b/db/intent_admission.go new file mode 100644 index 0000000..5c7b0dc --- /dev/null +++ b/db/intent_admission.go @@ -0,0 +1,32 @@ +package db + +import "fmt" + +// DiscardOpenIntent compensates a follow-up creation when task admission fails. +// The task execution gate must still be held by the caller, so the intent cannot +// be claimed between this check and deletion. Edges and anchors cascade with the +// node; activity is deleted explicitly because its node FK otherwise becomes NULL. +func (s *ExplorationStore) DiscardOpenIntent(id int64) error { + tx, err := s.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil { + return err + } + res, err := tx.Exec(`DELETE FROM exploration_nodes + WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID) + if err != nil { + return err + } + removed, err := res.RowsAffected() + if err != nil { + return err + } + if removed != 1 { + return fmt.Errorf("open intent %d was not available for admission rollback", id) + } + return tx.Commit() +} diff --git a/db/intent_control_test.go b/db/intent_control_test.go new file mode 100644 index 0000000..6cfb9cb --- /dev/null +++ b/db/intent_control_test.go @@ -0,0 +1,285 @@ +package db + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + expID, err := d.CreateExploration("intent CAS", "only one controller wins") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + store := d.Exploration(expID) + intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed { + t.Fatalf("claim: claimed=%v err=%v", claimed, err) + } + + const contenders = 12 + var winners atomic.Int32 + var wg sync.WaitGroup + start := make(chan struct{}) + for range contenders { + wg.Add(1) + go func() { + defer wg.Done() + <-start + changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused") + if transitionErr != nil { + t.Errorf("transition: %v", transitionErr) + return + } + if changed { + winners.Add(1) + } + }() + } + close(start) + wg.Wait() + if got := winners.Load(); got != 1 { + t.Fatalf("CAS winners=%d, want 1", got) + } + node, err := store.GetNode(intentID) + if err != nil || node == nil || node.State != "paused" { + t.Fatalf("node=%+v err=%v, want paused", node, err) + } +} + +func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + store := d.Exploration(expID) + intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed { + t.Fatalf("claim: claimed=%v err=%v", claimed, err) + } + + baselineDaily, err := d.TokenDailyAll(30) + if err != nil { + t.Fatal(err) + } + baseline := dailyTokenBuckets(baselineDaily) + + appendUsage := func(kind string, at time.Time, input, output, read, write *int) { + t.Helper() + activityID, appendErr := store.AppendActivity(Activity{ + NodeID: &intentID, Worker: "work#1", Kind: kind, + InputTokens: input, OutputTokens: output, + CacheReadTokens: read, CacheWriteTokens: write, + }) + if appendErr != nil { + t.Fatal(appendErr) + } + if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil { + t.Fatal(updateErr) + } + } + values := func(input, output, read, write int) (*int, *int, *int, *int) { + return &input, &output, &read, &write + } + + now := time.Now().UTC() + atUTCNoon := func(daysAgo int) time.Time { + at := now.AddDate(0, 0, -daysAgo) + return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC) + } + authoritativeDay := atUTCNoon(6) + fallbackDay := atUTCNoon(4) + trailingDay := atUTCNoon(2) + + // Run 1 has an authoritative result; its preceding cumulative frame must not + // be counted a second time or move usage to the frame's earlier date. + i, o, r, w := values(100, 100, 100, 100) + appendUsage("usage", atUTCNoon(7), i, o, r, w) + i, o, r, w = values(10, 11, 12, 13) + appendUsage("result", authoritativeDay, i, o, r, w) + // Run 2 models a legacy failed result that omitted usage. Its fallback belongs + // to the result's date, not the cumulative frame's date. + i, o, r, w = values(20, 21, 22, 23) + appendUsage("usage", atUTCNoon(5), i, o, r, w) + appendUsage("result", fallbackDay, nil, nil, nil, nil) + // Run 3 was interrupted before a terminal result was persisted. + i, o, r, w = values(30, 31, 32, 33) + appendUsage("usage", trailingDay, i, o, r, w) + + beforeCancelDaily, err := d.TokenDailyAll(30) + if err != nil { + t.Fatal(err) + } + beforeCancel := dailyTokenBuckets(beforeCancelDaily) + assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12) + assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0) + assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0) + + if err := store.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + cleanup, err := store.CancelIntent(intentID) + if err != nil { + t.Fatal(err) + } + if cleanup.Activities != 5 { + t.Fatalf("deleted activities=%d, want 5", cleanup.Activities) + } + total, err := store.TokenTotal() + if err != nil { + t.Fatal(err) + } + assertTokenUsage(t, total, 60, 63, 66, 69) + afterCancelDaily, err := d.TokenDailyAll(30) + if err != nil { + t.Fatal(err) + } + afterCancel := dailyTokenBuckets(afterCancelDaily) + assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12) + assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22) + assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32) + assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0) + + sessions, err := store.TokenStatsBySession() + if err != nil { + t.Fatal(err) + } + if len(sessions) != 0 { + t.Fatalf("cancelled intent leaked into sessions: %+v", sessions) + } + var rollups, datedRollups int + if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER ( + WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD') + ) FROM activity + WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result' + AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil { + t.Fatal(err) + } + if rollups != 3 || datedRollups != rollups { + t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups) + } + if _, err := store.CancelIntent(intentID); err == nil { + t.Fatal("second cancellation unexpectedly succeeded") + } + afterRetry, err := store.TokenTotal() + if err != nil { + t.Fatal(err) + } + assertTokenUsage(t, afterRetry, 60, 63, 66, 69) + afterRetryDaily, err := d.TokenDailyAll(30) + if err != nil { + t.Fatal(err) + } + assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel) +} + +func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket { + out := make(map[string]DailyTokenBucket, len(items)) + for _, item := range items { + out[item.Day] = item + } + return out +} + +func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) { + t.Helper() + day := at.UTC().Format(time.DateOnly) + actual, initial := got[day], baseline[day] + if actual.InputTokens-initial.InputTokens != input || + actual.OutputTokens-initial.OutputTokens != output || + actual.CacheReadTokens-initial.CacheReadTokens != read { + t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d", + day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens, + actual.CacheReadTokens-initial.CacheReadTokens, input, output, read) + } +} + +func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want) + } + for day, expected := range want { + if actual, ok := got[day]; !ok || actual != expected { + t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected) + } + } +} + +func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + store := d.Exploration(expID) + first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + for _, intentID := range []int64{first, second} { + if err := store.Link(intentID, RelYields, fact); err != nil { + t.Fatal(err) + } + if err := store.Link(intentID, RelYields, finding); err != nil { + t.Fatal(err) + } + } + if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed { + t.Fatalf("claim: claimed=%v err=%v", claimed, err) + } + if err := store.SetIntentState(first, "paused"); err != nil { + t.Fatal(err) + } + + cleanup, err := store.CancelIntent(first) + if err != nil { + t.Fatal(err) + } + if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 { + t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup) + } + for _, nodeID := range []int64{fact, finding} { + node, getErr := store.GetNode(nodeID) + if getErr != nil || node == nil { + t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr) + } + } +} diff --git a/db/task_assets_context.go b/db/task_assets_context.go index 2c5b1f2..b14c79a 100644 --- a/db/task_assets_context.go +++ b/db/task_assets_context.go @@ -29,8 +29,12 @@ target AS ( JOIN task_scope ts ON ( (ts.kind='company' AND a.company_id = ts.company_id) OR (ts.kind='root_domain' AND a.root_domain = ts.domain) - OR (ts.kind='subdomain' AND a.domain = ts.domain) - OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='subdomain' AND a.domain = ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='icp' AND ( + lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value + OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value + )) ) JOIN context_tasks ctx ON ctx.task_id=ts.task_id UNION @@ -53,7 +57,7 @@ tested AS ( func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) { rows, err := s.db.Query(`WITH `+directTaskContextCTE+` SELECT ts.id, ts.task_id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''), - COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'') + COALESCE(ts.net::text,''), COALESCE(ts.value,''), ts.source, COALESCE(ts.reason,'') FROM task_scope ts JOIN context_tasks ctx ON ctx.task_id=ts.task_id LEFT JOIN companies c ON c.id=ts.company_id @@ -66,7 +70,7 @@ ORDER BY CASE WHEN ts.task_id=$1 THEN 0 ELSE 1 END, ts.id`, taskID) for rows.Next() { var scope TaskScope var companyID int64 - if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Source, &scope.Reason); err != nil { + if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Value, &scope.Source, &scope.Reason); err != nil { return nil, err } if companyID > 0 { @@ -175,8 +179,12 @@ context_assets AS ( JOIN task_scope ts ON ( (ts.kind='company' AND a.company_id=ts.company_id) OR (ts.kind='root_domain' AND a.root_domain=ts.domain) - OR (ts.kind='subdomain' AND a.domain=ts.domain) - OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='subdomain' AND a.domain=ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='icp' AND ( + lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value + OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value + )) ) JOIN context_tasks ctx ON ctx.task_id=ts.task_id ) diff --git a/db/task_context.go b/db/task_context.go new file mode 100644 index 0000000..14053d4 --- /dev/null +++ b/db/task_context.go @@ -0,0 +1,419 @@ +package db + +import ( + "database/sql" + "fmt" + "strings" + "time" + "unicode/utf8" +) + +const IntentBlockedLLMQuota = "llm_quota_exhausted" + +// TaskLLMProfile is one ordered entry in a task's explicit failover chain. +type TaskLLMProfile struct { + ProfileID int64 `json:"profile_id"` + Position int `json:"position"` + Status string `json:"status"` + LastError string `json:"last_error,omitempty"` + ExhaustedAt *time.Time `json:"exhausted_at,omitempty"` +} + +// TaskSource identifies one directly related task and its exploration. +type TaskSource struct { + TaskID int64 + ExplorationID int64 + Description string + Goal string + Status string +} + +// TaskLLMTransition reports the shared task-level result of marking one profile +// quota-exhausted. NextProfileID is nil when the explicit chain is exhausted. +type TaskLLMTransition struct { + PreviousProfileID int64 + NextProfileID *int64 + ChainExhausted bool + Advanced bool + // Stale means the chain was replaced after the failing call selected its + // provider. The caller may retry a pre-stream request against the new chain, + // but must not report or persist a transition for this result. + Stale bool +} + +func (d *DB) hydrateTaskContext(t *Task) error { + if t == nil { + return nil + } + // A nil cursor on a non-empty chain is a persisted end-of-chain marker. Do not + // repair it from an earlier ready entry: the user may have manually selected a + // profile in the middle and legitimately exhausted every candidate after it. + legacyProfileID, activeProfileID, revision, chain, err := d.taskLLMContext(t.ID) + if err != nil { + return err + } + t.LLMProfileID = legacyProfileID + t.ActiveLLMProfileID = activeProfileID + t.LLMChainRevision = revision + sources, err := d.TaskSourceIDs(t.ID) + if err != nil { + return err + } + t.SourceTaskIDs = sources + companies, err := d.TaskCompanyIDs(t.ID) + if err != nil { + return err + } + t.CompanyIDs = companies + t.LLMProfileIDs = make([]int64, 0, len(chain)) + t.LLMFailoverState = "default" + var latest *TaskLLMProfile + activeReady := false + for i := range chain { + entry := chain[i] + t.LLMProfileIDs = append(t.LLMProfileIDs, entry.ProfileID) + if t.ActiveLLMProfileID != nil && entry.ProfileID == *t.ActiveLLMProfileID && entry.Status == "ready" { + activeReady = true + } + if entry.ExhaustedAt != nil && (latest == nil || entry.ExhaustedAt.After(*latest.ExhaustedAt)) { + copy := entry + latest = © + } + } + if len(chain) > 0 { + if activeReady { + t.LLMFailoverState = "ready" + } else { + t.LLMFailoverState = "chain_exhausted" + } + } + if latest != nil { + t.LLMFailoverReason = latest.LastError + } + return nil +} + +// taskLLMContext reads the compatibility profile, current cursor, and ordered +// chain in one statement. Runtime chain edits commit atomically, and one SQL +// statement gives hydration one matching snapshot instead of a transient mix of +// an old task cursor and a newly replaced chain (or vice versa). +func (d *DB) taskLLMContext(taskID int64) (*int64, *int64, int64, []TaskLLMProfile, error) { + rows, err := d.Query(` +SELECT t.llm_profile_id, t.active_llm_profile_id, t.llm_chain_revision, + p.profile_id, p.position, p.status, COALESCE(p.last_error,''), p.exhausted_at +FROM tasks t +LEFT JOIN task_llm_profiles p ON p.task_id=t.id +WHERE t.id=$1 +ORDER BY p.position NULLS LAST`, taskID) + if err != nil { + return nil, nil, 0, nil, err + } + defer rows.Close() + var ( + legacyProfileID *int64 + activeProfileID *int64 + revision int64 + chain []TaskLLMProfile + found bool + ) + for rows.Next() { + var legacy, active, profileID, position sql.NullInt64 + var rowRevision int64 + var status, lastError sql.NullString + var exhaustedAt sql.NullTime + if err := rows.Scan(&legacy, &active, &rowRevision, &profileID, &position, &status, &lastError, &exhaustedAt); err != nil { + return nil, nil, 0, nil, err + } + if !found { + found = true + revision = rowRevision + if legacy.Valid { + id := legacy.Int64 + legacyProfileID = &id + } + if active.Valid { + id := active.Int64 + activeProfileID = &id + } + } + if !profileID.Valid { + continue + } + entry := TaskLLMProfile{ + ProfileID: profileID.Int64, + Position: int(position.Int64), + Status: status.String, + LastError: lastError.String, + } + if exhaustedAt.Valid { + ts := exhaustedAt.Time + entry.ExhaustedAt = &ts + } + chain = append(chain, entry) + } + if err := rows.Err(); err != nil { + return nil, nil, 0, nil, err + } + if !found { + return nil, nil, 0, nil, sql.ErrNoRows + } + return legacyProfileID, activeProfileID, revision, chain, nil +} + +func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) { + rows, err := d.Query(`SELECT source_task_id FROM task_relations WHERE task_id=$1 ORDER BY created_at, source_task_id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + out = append(out, id) + } + return out, rows.Err() +} + +// TaskCompanyIDs returns the companies whose asset scopes are available to the +// task. The task_scope rows remain the single source of truth. +func (d *DB) TaskCompanyIDs(taskID int64) ([]int64, error) { + rows, err := d.Query(` +SELECT company_id FROM task_scope +WHERE task_id=$1 AND kind='company' AND company_id IS NOT NULL +ORDER BY id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + out = append(out, id) + } + return out, rows.Err() +} + +func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) { + rows, err := d.Query(` +SELECT t.id, t.exploration_id, t.description, t.goal, t.status +FROM task_relations r +JOIN tasks t ON t.id=r.source_task_id AND t.deleted_at IS NULL +WHERE r.task_id=$1 +ORDER BY r.created_at, r.source_task_id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []TaskSource + for rows.Next() { + var source TaskSource + if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil { + return nil, err + } + out = append(out, source) + } + return out, rows.Err() +} + +func (d *DB) TaskLLMProfiles(taskID int64) ([]TaskLLMProfile, error) { + rows, err := d.Query(`SELECT profile_id, position, status, COALESCE(last_error,''), exhausted_at +FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []TaskLLMProfile + for rows.Next() { + var entry TaskLLMProfile + if err := rows.Scan(&entry.ProfileID, &entry.Position, &entry.Status, &entry.LastError, &entry.ExhaustedAt); err != nil { + return nil, err + } + out = append(out, entry) + } + return out, rows.Err() +} + +// ReplaceTaskLLMProfiles atomically replaces and resets the explicit task chain. +// activeProfileID=0 selects the first entry. An empty list restores the existing +// agent-binding/global fallback behavior. +func (d *DB) ReplaceTaskLLMProfiles(taskID int64, profileIDs []int64, activeProfileID int64) error { + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + // Keep the task -> profile lock order shared by all task LLM mutations and + // DeleteProfile. Inserts below may take KEY SHARE locks on llm_profiles for + // their foreign keys, so the task row must be locked before any of them. + var status string + if err := tx.QueryRow(`SELECT status FROM tasks WHERE id=$1 AND deleted_at IS NULL FOR UPDATE`, taskID).Scan(&status); err != nil { + if err == sql.ErrNoRows { + return fmt.Errorf("task %d not found", taskID) + } + return err + } + if IsTerminal(status) { + return fmt.Errorf("terminal task cannot change LLM profiles") + } + if _, err := tx.Exec(`DELETE FROM task_llm_profiles WHERE task_id=$1`, taskID); err != nil { + return err + } + if err := insertTaskLLMProfiles(tx, taskID, profileIDs); err != nil { + return err + } + if len(profileIDs) == 0 { + activeProfileID = 0 + } else if activeProfileID == 0 { + activeProfileID = profileIDs[0] + } + if activeProfileID > 0 { + found := false + for _, id := range profileIDs { + found = found || id == activeProfileID + } + if !found { + return fmt.Errorf("active LLM profile %d is not in the task chain", activeProfileID) + } + } + var active any + if activeProfileID > 0 { + active = activeProfileID + } + if _, err := tx.Exec(`UPDATE tasks +SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, taskID, active); err != nil { + return err + } + return tx.Commit() +} + +// MarkTaskLLMProfileQuotaExhausted advances the shared task cursor once. A late +// in-flight error from an older profile records that entry as exhausted but does +// not advance past the profile another call already selected. +func (d *DB) MarkTaskLLMProfileQuotaExhausted(taskID, profileID int64, reason string) (TaskLLMTransition, error) { + return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, 0, false, reason) +} + +// MarkTaskLLMProfileQuotaExhaustedAtRevision applies a provider failure only +// when it belongs to the chain snapshot used to start that call. +func (d *DB) MarkTaskLLMProfileQuotaExhaustedAtRevision(taskID, profileID, revision int64, reason string) (TaskLLMTransition, error) { + return d.markTaskLLMProfileQuotaExhausted(taskID, profileID, revision, true, reason) +} + +func (d *DB) markTaskLLMProfileQuotaExhausted(taskID, profileID, revision int64, checkRevision bool, reason string) (TaskLLMTransition, error) { + var out TaskLLMTransition + tx, err := d.Begin() + if err != nil { + return out, err + } + defer tx.Rollback() + var ( + active sql.NullInt64 + currentRevision int64 + ) + if err := tx.QueryRow(`SELECT active_llm_profile_id, llm_chain_revision FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&active, ¤tRevision); err != nil { + return out, err + } + out.PreviousProfileID = profileID + if checkRevision && revision != currentRevision { + out.Stale = true + return out, tx.Commit() + } + var ( + position int + entryStatus string + ) + if err := tx.QueryRow(`SELECT position, status FROM task_llm_profiles WHERE task_id=$1 AND profile_id=$2`, taskID, profileID). + Scan(&position, &entryStatus); err != nil { + if err == sql.ErrNoRows { + // The chain was edited while this request was in flight. Its provider + // error remains valid for the caller, but it must not mutate the new chain. + return out, tx.Commit() + } + return out, err + } + // Once the cursor reached the end, late failures must be idempotent. This also + // prevents an older in-flight request from reviving a ready entry before a + // manually selected starting position. + if !active.Valid { + out.ChainExhausted = true + return out, tx.Commit() + } + reason = truncateUTF8(strings.TrimSpace(reason), 1000) + if entryStatus != "quota_exhausted" { + if _, err := tx.Exec(`UPDATE task_llm_profiles +SET status='quota_exhausted', last_error=$3, exhausted_at=now() +WHERE task_id=$1 AND profile_id=$2`, taskID, profileID, reason); err != nil { + return out, err + } + } + if active.Valid && active.Int64 != profileID { + next := active.Int64 + out.NextProfileID = &next + return out, tx.Commit() + } + var next int64 + err = tx.QueryRow(`SELECT profile_id FROM task_llm_profiles +WHERE task_id=$1 AND position>$2 AND status='ready' +ORDER BY position LIMIT 1`, taskID, position).Scan(&next) + switch err { + case nil: + out.Advanced = true + out.NextProfileID = &next + if _, err := tx.Exec(`UPDATE tasks + SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1 + WHERE id=$1`, taskID, next); err != nil { + return out, err + } + case sql.ErrNoRows: + out.Advanced = true + out.ChainExhausted = true + if _, err := tx.Exec(`UPDATE tasks + SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1 + WHERE id=$1`, taskID); err != nil { + return out, err + } + default: + return out, err + } + return out, tx.Commit() +} + +func truncateUTF8(value string, maxBytes int) string { + value = strings.ToValidUTF8(value, "\uFFFD") + if maxBytes <= 0 { + return "" + } + if len(value) <= maxBytes { + return value + } + end := maxBytes + for end > 0 && !utf8.ValidString(value[:end]) { + end-- + } + return value[:end] +} + +func (s *ExplorationStore) SetIntentBlockedReason(id int64, reason string) error { + _, err := s.db.Exec(`UPDATE exploration_nodes +SET state='blocked', blocked_reason=NULLIF($1,''), completed_at=now() +WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, reason, id, s.expID) + return err +} + +func (s *ExplorationStore) ReopenIntentsByBlockedReason(reason string) (int64, error) { + res, err := s.db.Exec(`UPDATE exploration_nodes +SET state='open', blocked_reason=NULL, completed_at=NULL +WHERE exploration_id=$1 AND kind='intent' AND state='blocked' AND blocked_reason=$2`, s.expID, reason) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + return n, nil +} diff --git a/db/task_context_lock_test.go b/db/task_context_lock_test.go new file mode 100644 index 0000000..06b455d --- /dev/null +++ b/db/task_context_lock_test.go @@ -0,0 +1,527 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" +) + +func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + first, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-first-%d", suffix), Format: "openai", Model: "first", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + second, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-second-%d", suffix), Format: "openai", Model: "second", APIKey: "test-key", + }) + if err != nil { + _ = d.DeleteProfile(first) + t.Fatal(err) + } + task, err := d.CreateTaskWithOptions("LLM lock order", "verify concurrent mutation locks", TaskCreateOptions{ + LLMProfileIDs: []int64{first, second}, + }) + if err != nil { + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + }) + + replaceDB, replacePID := openSingleConnectionTestDB(t, dsn) + replaceBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer replaceBlocker.Rollback() + if _, err := replaceBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + replaceDone := make(chan error, 1) + go func() { + replaceDone <- replaceDB.ReplaceTaskLLMProfiles(task.ID, []int64{second, first}, second) + }() + if err := waitForBackendBlock(d, replacePID, replaceDone); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, first, second) + if err := replaceBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(replaceDone); err != nil { + t.Fatalf("replace chain: %v", err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer deleteBlocker.Rollback() + if _, err := deleteBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + deleteDone := make(chan error, 1) + go func() { + deleteDone <- deleteDB.DeleteProfile(second) + }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, second) + if err := deleteBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != first { + t.Fatalf("deleting the active profile did not select its successor: %+v", got) + } +} + +func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-lock-profile-%d", suffix), Format: "openai", Model: "agent-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("lock_agent_%d", suffix), "lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + if err := d.SetAgentLLMProfile(agent.Key, &profileID); err != nil { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete agent profile: %v", err) + } + got, err := d.GetAgentByKey(agent.Key) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("agent binding was not cleared: agent=%+v err=%v", got, err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-lock-profile-%d", suffix), Format: "openai", Model: "conversation-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "lock conversation", &profileID) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete conversation profile: %v", err) + } + got, err := d.GetConversation(conversation.ID) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("conversation binding was not cleared: conversation=%+v err=%v", got, err) + } + }) +} + +func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-write-lock-profile-%d", suffix), Format: "openai", Model: "agent-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("write_lock_agent_%d", suffix), "write lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.SetAgentLLMProfile(agent.Key, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind agent profile: %v", err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-write-lock-profile-%d", suffix), Format: "openai", Model: "conversation-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "write lock conversation", nil) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.UpdateConversationProfile(conversation.ID, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind conversation profile: %v", err) + } + }) +} + +func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-create-race-profile-%d", suffix), Format: "openai", Model: "conversation-create-race", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + title := fmt.Sprintf("conversation create race %d", suffix) + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM conversations WHERE title=$1`, title) + _ = d.DeleteProfile(profileID) + }) + + profileBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer profileBlocker.Rollback() + if _, err := profileBlocker.Exec(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE`, profileID); err != nil { + t.Fatal(err) + } + + createDB, createPID := openSingleConnectionTestDB(t, dsn) + type createResult struct { + conversation *Conversation + err error + } + createDone := make(chan createResult, 1) + createStatus := make(chan error, 1) + go func() { + conversation, err := createDB.CreateConversation("planner", title, &profileID) + createDone <- createResult{conversation: conversation, err: err} + createStatus <- err + }() + if err := waitForBackendBlock(d, createPID, createStatus); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + if err := profileBlocker.Rollback(); err != nil { + t.Fatal(err) + } + + var created createResult + select { + case created = <-createDone: + case <-time.After(12 * time.Second): + t.Fatal("timed out waiting for conversation creation") + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile during conversation creation: %v", err) + } + var persisted int + if err := d.QueryRow(`SELECT count(*) FROM conversations WHERE title=$1`, title).Scan(&persisted); err != nil { + t.Fatal(err) + } + if created.err == nil { + if created.conversation == nil || persisted != 1 { + t.Fatalf("successful creation was not committed atomically: conversation=%+v count=%d", created.conversation, persisted) + } + } else if persisted != 0 { + t.Fatalf("failed creation left a partial conversation row: err=%v count=%d", created.err, persisted) + } +} + +func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("late-reference-profile-%d", suffix), Format: "openai", Model: "late-reference", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + task, err := d.CreateTask("late profile reference", "exercise delete retry", nil, 0, 0) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(profileID) + }) + + // Keep the new reference uncommitted while deletion takes its initial + // READ COMMITTED snapshot. The task is therefore absent from the first lock + // set, while its FK KEY SHARE lock makes deletion wait at the profile row. + referenceTx, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer referenceTx.Rollback() + if _, err := referenceTx.Exec(`UPDATE tasks +SET llm_profile_id=$2, active_llm_profile_id=$2 +WHERE id=$1`, task.ID, profileID); err != nil { + t.Fatal(err) + } + if _, err := referenceTx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position) +VALUES ($1,$2,0)`, task.ID, profileID); err != nil { + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + referenceTx.Rollback() + t.Fatal(err) + } + if err := referenceTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile after late reference: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMChainRevision != 1 { + t.Fatalf("late task reference was not handled by a locked retry: %+v", got) + } +} + +func openSingleConnectionTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForBackendBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("mutation returned before reaching the expected reference-row lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func assertProfilesUnlocked(t *testing.T, d *DB, profileIDs ...int64) { + t.Helper() + probe, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer probe.Rollback() + for _, profileID := range profileIDs { + var lockedID int64 + if err := probe.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE NOWAIT`, profileID).Scan(&lockedID); err != nil { + t.Fatalf("profile %d was locked before the task row: %v", profileID, err) + } + } +} + +func waitForMutationResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for task LLM mutation") + } +} diff --git a/db/task_context_unit_test.go b/db/task_context_unit_test.go new file mode 100644 index 0000000..519f304 --- /dev/null +++ b/db/task_context_unit_test.go @@ -0,0 +1,27 @@ +package db + +import ( + "strings" + "testing" + "unicode/utf8" +) + +func TestTruncateUTF8PreservesValidEncoding(t *testing.T) { + t.Parallel() + input := strings.Repeat("额度不足", 400) + got := truncateUTF8(input, 1000) + if len(got) > 1000 { + t.Fatalf("truncated value has %d bytes, want at most 1000", len(got)) + } + if !utf8.ValidString(got) { + t.Fatalf("truncated value is not valid UTF-8: %q", got[len(got)-8:]) + } +} + +func TestTruncateUTF8RepairsInvalidInput(t *testing.T) { + t.Parallel() + got := truncateUTF8("bad\xffvalue", 1000) + if !utf8.ValidString(got) { + t.Fatalf("repaired value is not valid UTF-8: %q", got) + } +} diff --git a/db/task_delete_concurrency_test.go b/db/task_delete_concurrency_test.go new file mode 100644 index 0000000..0d65191 --- /dev/null +++ b/db/task_delete_concurrency_test.go @@ -0,0 +1,210 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" + + _ "github.com/jackc/pgx/v5/stdlib" +) + +func TestDeleteTaskTrafficHostsUseOneLockedTransaction(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + t.Run("sharing committed before deletion lock is observed", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/concurrent-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + writer, _ := openTaskDeleteTestDB(t, dsn) + writerTx, err := writer.Begin() + if err != nil { + t.Fatal(err) + } + defer writerTx.Rollback() + if _, err := writerTx.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID); err != nil { + t.Fatal(err) + } + + deleter, deleterPID := openTaskDeleteTestDB(t, dsn) + var preparedHosts []string + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + preparedHosts = append([]string(nil), p.TrafficHosts...) + return nil + }) + deleteDone <- deleteErr + }() + + if err := waitForTaskDeleteBlock(d, deleterPID, deleteDone); err != nil { + t.Fatal(err) + } + if err := writerTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if containsDeleteHost(preparedHosts, host) { + t.Fatalf("newly shared host %q was selected for traffic deletion: %v", host, preparedHosts) + } + var remaining int + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE url=$1 AND $2=ANY(task_ids)`, serviceURL, second.ID).Scan(&remaining); err != nil { + t.Fatal(err) + } + if remaining != 1 { + t.Fatalf("concurrent owner's asset was not preserved: count=%d", remaining) + } + }) + + t.Run("sharing cannot enter between host resolution and commit", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/late-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + deleter, _ := openTaskDeleteTestDB(t, dsn) + prepared := make(chan TaskDeletePreparation, 1) + releasePrepare := make(chan struct{}) + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + prepared <- p + <-releasePrepare + return nil + }) + deleteDone <- deleteErr + }() + + var plan TaskDeletePreparation + select { + case plan = <-prepared: + case err := <-deleteDone: + close(releasePrepare) + t.Fatalf("delete returned before preparation: %v", err) + case <-time.After(5 * time.Second): + close(releasePrepare) + t.Fatal("timed out waiting for task delete preparation") + } + if !containsDeleteHost(plan.TrafficHosts, host) { + close(releasePrepare) + t.Fatalf("exclusive host %q missing from preparation: %v", host, plan.TrafficHosts) + } + + writer, writerPID := openTaskDeleteTestDB(t, dsn) + writerDone := make(chan error, 1) + go func() { + _, writeErr := writer.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID) + writerDone <- writeErr + }() + if err := waitForTaskDeleteBlock(d, writerPID, writerDone); err != nil { + close(releasePrepare) + t.Fatal(err) + } + + close(releasePrepare) + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(writerDone); err != nil { + t.Fatal(err) + } + }) +} + +func createTaskDeleteRaceFixture(t *testing.T, d *DB) (first, second *Task, host string, rootAssetID int64) { + t.Helper() + var err error + first, err = d.CreateTask("delete race owner", "delete safely", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err = d.CreateTask("delete race sharer", "preserve shared host", nil, 0, 0) + if err != nil { + _ = d.DeleteTask(first.ID) + t.Fatal(err) + } + host = fmt.Sprintf("task-delete-race-%d.example.test", first.ID) + rootAssetID, err = d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: host, TaskID: first.ID}) + if err != nil { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + t.Fatal(err) + } + return first, second, host, rootAssetID +} + +func openTaskDeleteTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForTaskDeleteBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("operation returned before reaching the task deletion lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func waitForTaskDeleteResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for concurrent task deletion operation") + } +} + +func containsDeleteHost(hosts []string, want string) bool { + for _, host := range hosts { + if host == want { + return true + } + } + return false +} diff --git a/db/task_queue_test.go b/db/task_queue_test.go new file mode 100644 index 0000000..a9b949b --- /dev/null +++ b/db/task_queue_test.go @@ -0,0 +1,66 @@ +package db + +import ( + "testing" + "time" +) + +func TestTaskQueuePreservesBootstrapAndFIFOPosition(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + task, err := d.CreateTask("queue metadata", "keep first-run mode", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer d.DeleteTask(task.ID) + + if err := d.Enqueue(task.ID, "bootstrap"); err != nil { + t.Fatal(err) + } + first, err := d.GetTask(task.ID) + if err != nil || first == nil || first.QueuedAt == nil { + t.Fatalf("first enqueue: task=%+v err=%v", first, err) + } + + // A follow-up/rerun can try to admit an already queued task as resume. The + // original bootstrap mode and FIFO timestamp must remain authoritative. + time.Sleep(time.Millisecond) + if err := d.Enqueue(task.ID, "resume"); err != nil { + t.Fatal(err) + } + second, err := d.GetTask(task.ID) + if err != nil || second == nil { + t.Fatalf("second enqueue: task=%+v err=%v", second, err) + } + if second.QueueMode != "bootstrap" { + t.Fatalf("queue mode=%q, want bootstrap", second.QueueMode) + } + if second.QueuedAt == nil || !second.QueuedAt.Equal(*first.QueuedAt) { + t.Fatalf("repeated enqueue moved FIFO position: first=%v second=%v", first.QueuedAt, second.QueuedAt) + } + + // Pausing a queued task removes it from the queue but keeps its required + // startup mode. A later requeue receives a new tail position. + if err := d.Dequeue(task.ID, false); err != nil { + t.Fatal(err) + } + paused, err := d.GetTask(task.ID) + if err != nil || paused == nil || paused.Queued || paused.QueueMode != "bootstrap" { + t.Fatalf("paused queue metadata: task=%+v err=%v", paused, err) + } + time.Sleep(time.Millisecond) + if err := d.Enqueue(task.ID, "bootstrap"); err != nil { + t.Fatal(err) + } + requeued, err := d.GetTask(task.ID) + if err != nil || requeued == nil || requeued.QueuedAt == nil { + t.Fatalf("requeue: task=%+v err=%v", requeued, err) + } + if !requeued.QueuedAt.After(*first.QueuedAt) { + t.Fatalf("requeue did not move to FIFO tail: first=%v requeued=%v", first.QueuedAt, requeued.QueuedAt) + } +} diff --git a/db/task_source_limit_test.go b/db/task_source_limit_test.go new file mode 100644 index 0000000..7c54022 --- /dev/null +++ b/db/task_source_limit_test.go @@ -0,0 +1,47 @@ +package db + +import ( + "errors" + "reflect" + "strings" + "testing" +) + +func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) { + sourceIDs := make([]int64, MaxTaskSourceCount+1) + for i := range sourceIDs { + sourceIDs[i] = int64(i + 1) + } + + // No database handle is needed: validation must run before Begin so an + // oversized request cannot consume a connection or create partial rows. + _, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{SourceTaskIDs: sourceIDs}) + if err == nil || !strings.Contains(err.Error(), "too many source tasks") { + t.Fatalf("expected source-count validation error, got %v", err) + } +} + +func TestNormalizeTaskCompanyIDs(t *testing.T) { + got, err := NormalizeTaskCompanyIDs([]int64{4, 2, 4, 7, 2}) + if err != nil { + t.Fatal(err) + } + if want := []int64{4, 2, 7}; !reflect.DeepEqual(got, want) { + t.Fatalf("NormalizeTaskCompanyIDs=%v, want %v", got, want) + } + if _, err := NormalizeTaskCompanyIDs([]int64{1, 0}); !errors.Is(err, ErrTaskCompanyIDsInvalid) { + t.Fatalf("invalid company id error=%v", err) + } +} + +func TestCreateTaskRejectsTooManyCompaniesBeforeOpeningTransaction(t *testing.T) { + companyIDs := make([]int64, MaxTaskCompanyCount+1) + for i := range companyIDs { + companyIDs[i] = int64(i + 1) + } + + _, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{CompanyIDs: companyIDs}) + if !errors.Is(err, ErrTaskCompanyIDsInvalid) { + t.Fatalf("expected company-count validation error, got %v", err) + } +} diff --git a/db/tasks.go b/db/tasks.go index 1ce8c98..bc0a57b 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -1,7 +1,11 @@ package db import ( + "database/sql" "encoding/json" + "errors" + "fmt" + "strconv" "time" ) @@ -9,15 +13,29 @@ import ( type Task struct { ID int64 `json:"id"` Name string `json:"name"` // 可选任务名称;空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` + CategoryName string `json:"category_name,omitempty"` Description string `json:"description"` Goal string `json:"goal"` ExplorationID int64 `json:"exploration_id"` Status string `json:"status"` Paused bool `json:"paused"` + Queued bool `json:"queued"` + QueuedAt *time.Time `json:"queued_at,omitempty"` + QueueMode string `json:"queue_mode,omitempty"` LLMProfileID *int64 `json:"llm_profile_id,omitempty"` - ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) - CreatedAt time.Time `json:"created_at"` - CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil + // Task-level ordered LLM chain. LLMProfileID remains the compatibility alias + // for ActiveLLMProfileID while older API clients still send one profile id. + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMChainRevision int64 `json:"-"` + LLMFailoverState string `json:"llm_failover_state,omitempty"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` + CompanyIDs []int64 `json:"company_ids,omitempty"` + ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) + CreatedAt time.Time `json:"created_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil // 任务级超时(见 docs/任务级超时与收尾设计.md)。 TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时 FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行 @@ -32,6 +50,23 @@ type Task struct { CoverageEnabled bool `json:"coverage_enabled"` } +// TaskDeleteResult reports optional related-data cleanup performed in the same +// transaction as the task/exploration delete. +type TaskDeleteResult struct { + AssetsDeleted int64 + AssetsDetached int64 + FindingsDeleted int64 + LLMRecordsDeleted int64 +} + +// TaskDeletePreparation is produced inside the PostgreSQL deletion transaction +// after asset and anchor writers have been excluded. Prepare callbacks may use +// TrafficHosts to stage an external traffic deletion before PostgreSQL commits. +type TaskDeletePreparation struct { + ExplorationID int64 + TrafficHosts []string +} + // IsTerminal reports whether a task status is a terminal (finished) state. // 单一真源，替换散落各处的 done/failed 硬编码判定。 func IsTerminal(status string) bool { @@ -45,6 +80,46 @@ func IsTerminal(status string) bool { // 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min，防止把 planner 打爆。 const MinPlanHeartbeatSeconds = 600 +// MaxTaskSourceCount bounds the amount of live inherited context one task can +// pull into every planner/main-agent prompt. Inheritance is intentionally direct +// only; keeping the fan-in bounded also prevents a single create request from +// multiplying graph and asset-context queries without limit. +const MaxTaskSourceCount = 8 + +// MaxTaskCompanyCount bounds the number of company asset scopes attached to a +// task. Company scopes are prompt context, so an unbounded list would inflate +// every planner and worker call. +const MaxTaskCompanyCount = 32 + +var ( + ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids") + ErrTaskCompanyNotFound = errors.New("task company not found") +) + +// NormalizeTaskCompanyIDs validates IDs and removes duplicates while retaining +// the user's first-seen order. +func NormalizeTaskCompanyIDs(ids []int64) ([]int64, error) { + if len(ids) == 0 { + return nil, nil + } + seen := make(map[int64]struct{}, min(len(ids), MaxTaskCompanyCount)) + normalized := make([]int64, 0, min(len(ids), MaxTaskCompanyCount)) + for _, id := range ids { + if id <= 0 { + return nil, fmt.Errorf("%w: company id must be positive", ErrTaskCompanyIDsInvalid) + } + if _, exists := seen[id]; exists { + continue + } + seen[id] = struct{}{} + normalized = append(normalized, id) + if len(normalized) > MaxTaskCompanyCount { + return nil, fmt.Errorf("%w: got more than %d unique companies", ErrTaskCompanyIDsInvalid, MaxTaskCompanyCount) + } + } + return normalized, nil +} + func normalizeHeartbeat(sec int) int { if sec < MinPlanHeartbeatSeconds { return MinPlanHeartbeatSeconds @@ -53,8 +128,6 @@ func normalizeHeartbeat(sec int) int { } func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { -<<<<<<< Updated upstream -======= var ids []int64 if llmProfileID != nil { ids = []int64{*llmProfileID} @@ -68,6 +141,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe // with the task/exploration row. type TaskCreateOptions struct { Name string // 可选任务名称;空=未命名 + CategoryID *int64 SourceTaskIDs []int64 CompanyIDs []int64 LLMProfileIDs []int64 @@ -89,7 +163,6 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti return nil, err } opts.CompanyIDs = companyIDs ->>>>>>> Stashed changes tx, err := d.Begin() if err != nil { return nil, err @@ -114,22 +187,30 @@ INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, or VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil { return nil, err } -<<<<<<< Updated upstream - if timeoutSeconds < 0 { - timeoutSeconds = 0 -======= if opts.TimeoutSeconds < 0 { opts.TimeoutSeconds = 0 } opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds) coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled + var categoryName string + if opts.CategoryID != nil { + if *opts.CategoryID <= 0 { + return nil, fmt.Errorf("%w: category id must be positive", ErrTaskCategoryInvalid) + } + if err := tx.QueryRow(`SELECT name FROM task_categories WHERE id=$1`, *opts.CategoryID).Scan(&categoryName); err != nil { + if err == sql.ErrNoRows { + return nil, ErrTaskCategoryNotFound + } + return nil, err + } + } var active *int64 if len(opts.LLMProfileIDs) > 0 { id := opts.LLMProfileIDs[0] active = &id } t := &Task{ - Name: opts.Name, + Name: opts.Name, CategoryID: opts.CategoryID, CategoryName: categoryName, Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: active, ActiveLLMProfileID: active, LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...), @@ -137,31 +218,30 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); CompanyIDs: append([]int64(nil), opts.CompanyIDs...), TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds, CoverageEnabled: coverageEnabled, ->>>>>>> Stashed changes } - planHeartbeatSeconds = normalizeHeartbeat(planHeartbeatSeconds) // <=0→300，<30→30 - t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds} if err := tx.QueryRow(` -<<<<<<< Updated upstream -INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6) -RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { -======= -INSERT INTO tasks(name, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled) -VALUES ($1,$2,$3,$4,$5,$5,$6,$7,$8) -RETURNING id, status, paused, created_at`, opts.Name, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { ->>>>>>> Stashed changes +INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled) +VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9) +RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } + if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { + return nil, err + } + if err := insertTaskCompanies(tx, t.ID, opts.CompanyIDs); err != nil { + return nil, err + } + if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil { + return nil, err + } + if len(opts.LLMProfileIDs) == 0 { + t.LLMFailoverState = "default" + } else { + t.LLMFailoverState = "ready" + } return t, tx.Commit() } -<<<<<<< Updated upstream -const taskCols = `id, description, goal, exploration_id, status, paused, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at` - -func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { - var t Task - if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil { -======= func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error { if len(companyIDs) == 0 { return nil @@ -222,12 +302,13 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { return nil } -const taskCols = `id, COALESCE(name,''), description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at` +const taskCols = `id, COALESCE(name,''), category_id, +COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''), +description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Name, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { ->>>>>>> Stashed changes + if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } return &t, nil @@ -254,7 +335,16 @@ func (d *DB) ListTasks() ([]*Task, error) { } out = append(out, t) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + rows.Close() + for _, task := range out { + if err := d.hydrateTaskContext(task); err != nil { + return nil, err + } + } + return out, nil } // GetTask returns one alive task (nil if not found/deleted). @@ -266,6 +356,9 @@ func (d *DB) GetTask(id int64) (*Task, error) { } return nil, err } + if err := d.hydrateTaskContext(t); err != nil { + return nil, err + } return t, nil } @@ -275,6 +368,43 @@ func (d *DB) SetPaused(id int64, paused bool) error { return err } +// Enqueue places a task at the tail of the persistent admission queue. Repeating +// the operation while it is already queued keeps its original FIFO position. +func (d *DB) Enqueue(id int64, mode string) error { + if mode != "bootstrap" && mode != "resume" { + return fmt.Errorf("invalid queue mode %q", mode) + } + _, err := d.Exec(`UPDATE tasks +SET queued=true, + queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END, + queue_mode=CASE + WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap' + ELSE 'resume' + END +WHERE id=$1`, id, mode) + return err +} + +// Dequeue removes the concurrency hold. clearMode is false when a user pauses a +// queued task, preserving whether its next admission must bootstrap or resume. +func (d *DB) Dequeue(id int64, clearMode bool) error { + _, err := d.Exec(`UPDATE tasks +SET queued=false, + queued_at=NULL, + queue_mode=CASE WHEN $2 THEN '' ELSE queue_mode END +WHERE id=$1`, id, clearMode) + return err +} + +// SetQueued is the compatibility helper used by older callers and tests. New +// scheduling code should use Enqueue/Dequeue so FIFO metadata is explicit. +func (d *DB) SetQueued(id int64, queued bool) error { + if queued { + return d.Enqueue(id, "bootstrap") + } + return d.Dequeue(id, true) +} + // SetStatus updates a task's lifecycle status. Entering a terminal state // (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp // stable); moving back to a non-terminal state clears it, so a re-run has no stale @@ -323,28 +453,134 @@ UPDATE tasks return deadline, err } -// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/ -// activity via ON DELETE CASCADE). The task row is removed first so the -// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete. -// Global assets are NOT touched. +// DeleteTask preserves the historical behavior: remove the task and its +// exploration subgraph while retaining global assets. func (d *DB) DeleteTask(id int64) error { + _, err := d.DeleteTaskCascade(id, false, false) + return err +} + +// DeleteTaskCascade hard-deletes a task and its exploration subgraph +// (nodes/edges/activity via ON DELETE CASCADE). Optional standalone findings and +// assets owned only by this task are deleted in the same transaction; shared +// assets are retained and only have this task id detached. deleteLLMRecords is +// optional to preserve callers of the original two-option API. +func (d *DB) DeleteTaskCascade(id int64, deleteAssets, deleteFindings bool, deleteLLMRecords ...bool) (TaskDeleteResult, error) { + deleteRecords := len(deleteLLMRecords) > 0 && deleteLLMRecords[0] + return d.DeleteTaskCascadePrepared(id, deleteAssets, deleteFindings, deleteRecords, nil) +} + +// DeleteTaskCascadePrepared coordinates reversible external deletion with the +// PostgreSQL cascade. When prepare is non-nil, host ownership is resolved and +// prepare is invoked inside this transaction while asset and anchor writers are +// excluded. The locks remain held through commit, closing the window where a +// host could become shared after its traffic had already been staged. +// +// prepare must only stage reversible work. Any returned error or later database +// error rolls PostgreSQL back; the caller remains responsible for rolling back +// external work that its callback staged successfully. +func (d *DB) DeleteTaskCascadePrepared( + id int64, + deleteAssets, deleteFindings, deleteLLMRecords bool, + prepare func(TaskDeletePreparation) error, +) (TaskDeleteResult, error) { + var result TaskDeleteResult tx, err := d.Begin() if err != nil { - return err + return result, err } defer tx.Rollback() var expID int64 - if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil { - if err.Error() == "sql: no rows in result set" { - return nil // already gone + if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 FOR UPDATE`, id).Scan(&expID); err != nil { + if err == sql.ErrNoRows { + return result, nil // already gone } - return err + return result, err } + + // SHARE ROW EXCLUSIVE conflicts with every INSERT/UPDATE/DELETE on these + // tables and with another coordinated deletion. Taking both in one fixed order + // prevents phantoms (a distinct asset row for the same host) as well as new + // ownership/anchor references until the deletion transaction commits. + if deleteAssets || prepare != nil { + if _, err := tx.Exec(`LOCK TABLE assets, exploration_anchors IN SHARE ROW EXCLUSIVE MODE`); err != nil { + return result, err + } + } + if prepare != nil { + hosts, err := hostsForTaskDeletion(tx, id, expID) + if err != nil { + return result, err + } + if err := prepare(TaskDeletePreparation{ + ExplorationID: expID, + TrafficHosts: hosts, + }); err != nil { + return result, err + } + } + if deleteAssets { + // Ownership is the union of explicit task_ids and this exploration's anchors. + // An asset is deletable only when no other live task references it through + // either mechanism. This covers legacy anchor-only seeds without destroying + // evidence anchored by another task. + res, err := tx.Exec(` +WITH candidate_assets AS ( + SELECT id FROM assets WHERE $1 = ANY(task_ids) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + WHERE n.exploration_id=$2 +), +deletable AS ( + SELECT a.id + FROM assets a + JOIN candidate_assets c ON c.id=a.id + WHERE NOT EXISTS ( + SELECT 1 FROM tasks t + WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids) + ) AND NOT EXISTS ( + SELECT 1 + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + JOIN tasks t ON t.exploration_id=n.exploration_id + WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL + ) +) +DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID) + if err != nil { + return result, err + } + result.AssetsDeleted, _ = res.RowsAffected() + + res, err = tx.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, id) + if err != nil { + return result, err + } + result.AssetsDetached, _ = res.RowsAffected() + } + if deleteFindings { + res, err := tx.Exec(`DELETE FROM findings WHERE task_id = $1`, id) + if err != nil { + return result, err + } + result.FindingsDeleted, _ = res.RowsAffected() + } + if deleteLLMRecords { + res, err := tx.Exec(`DELETE FROM llm_records WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(id, 10)) + if err != nil { + return result, err + } + result.LLMRecordsDeleted, _ = res.RowsAffected() + } + // llm_usage (the token metering ledger) is intentionally NOT deleted with the + // task — it is kept as historical accounting even after the task is gone. if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil { - return err + return result, err } if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil { - return err + return result, err } - return tx.Commit() + return result, tx.Commit() } diff --git a/db/tasks_test.go b/db/tasks_test.go index f2e9def..0f622e3 100644 --- a/db/tasks_test.go +++ b/db/tasks_test.go @@ -1,6 +1,12 @@ package db -import "testing" +import ( + "errors" + "fmt" + "sync" + "testing" + "time" +) func TestTaskLifecycleAndDeleteCascade(t *testing.T) { d, err := Open(testDSN(t)) @@ -27,6 +33,23 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) { if err != nil || got == nil || !got.Paused { t.Fatalf("paused not persisted: %+v err=%v", got, err) } + if got.Queued { + t.Fatalf("new task should not be queued: %+v", got) + } + + // queued (concurrency-hold) flag round-trips independently of paused + if err := d.SetQueued(tk.ID, true); err != nil { + t.Fatal(err) + } + if g, _ := d.GetTask(tk.ID); g == nil || !g.Queued { + t.Fatalf("queued not persisted: %+v", g) + } + if err := d.SetQueued(tk.ID, false); err != nil { + t.Fatal(err) + } + if g, _ := d.GetTask(tk.ID); g == nil || g.Queued { + t.Fatalf("queued not cleared: %+v", g) + } // list contains it list, _ := d.ListTasks() @@ -58,3 +81,485 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) { t.Fatalf("exploration should be deleted, got %d", exps) } } + +func TestTaskDeleteCascadeAssets(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + first, err := d.CreateTask("级联删除测试", "目标A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err := d.CreateTask("共享资产保留测试", "目标B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + exclusiveID, err := assets.UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("delete-%d.example.test", first.ID), + TaskID: first.ID, + }) + if err != nil { + t.Fatal(err) + } + sharedDomain := fmt.Sprintf("shared-%d.example.test", first.ID) + sharedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + if _, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: second.ID}); err != nil { + t.Fatal(err) + } + anchorOnlyDomain := fmt.Sprintf("anchor-only-%d.example.test", first.ID) + anchorOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: anchorOnlyDomain}) + if err != nil { + t.Fatal(err) + } + firstOrigin, err := d.Exploration(first.ExplorationID).OriginFactID() + if err != nil || firstOrigin == 0 { + t.Fatalf("first origin: id=%d err=%v", firstOrigin, err) + } + if err := d.Exploration(first.ExplorationID).Anchor(firstOrigin, anchorOnlyID); err != nil { + t.Fatal(err) + } + protectedDomain := fmt.Sprintf("other-anchor-%d.example.test", first.ID) + protectedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: protectedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + secondOrigin, err := d.Exploration(second.ExplorationID).OriginFactID() + if err != nil || secondOrigin == 0 { + t.Fatalf("second origin: id=%d err=%v", secondOrigin, err) + } + if err := d.Exploration(second.ExplorationID).Anchor(secondOrigin, protectedID); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + }) + + hosts, err := assets.HostsByTask(first.ID) + if err != nil { + t.Fatal(err) + } + hostSet := make(map[string]bool, len(hosts)) + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !hostSet[sharedDomain] { + t.Fatalf("task hosts missing cascade fixtures: %v", hosts) + } + deletableHosts, err := assets.HostsForTaskDeletion(first.ID, first.ExplorationID) + if err != nil { + t.Fatal(err) + } + deletableSet := make(map[string]bool, len(deletableHosts)) + for _, host := range deletableHosts { + deletableSet[host] = true + } + if !deletableSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !deletableSet[anchorOnlyDomain] { + t.Fatalf("exclusive or anchor-only cleanup host missing: %v", deletableHosts) + } + if deletableSet[sharedDomain] || deletableSet[protectedDomain] { + t.Fatalf("shared host was not protected from traffic cleanup: %v", deletableHosts) + } + + findingID, err := d.AddFinding(first.ID, 0, "__task_delete_cascade__", "", SeverityHigh, "summary", "evidence", "tester", []int64{exclusiveID}) + if err != nil { + t.Fatal(err) + } + + result, err := d.DeleteTaskCascade(first.ID, true, true) + if err != nil { + t.Fatal(err) + } + if result.AssetsDeleted < 2 || result.AssetsDetached < 2 { + t.Fatalf("unexpected asset cleanup result: %+v", result) + } + if result.FindingsDeleted != 1 { + t.Fatalf("unexpected finding cleanup result: %+v", result) + } + if finding, err := d.GetFinding(findingID); err != nil || finding != nil { + t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err) + } + remaining, err := assets.GetByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + if err != nil { + t.Fatal(err) + } + if len(remaining) != 2 { + t.Fatalf("only shared and other-anchored assets should remain, got %+v", remaining) + } + byID := make(map[int64]*Asset, len(remaining)) + for _, asset := range remaining { + byID[asset.ID] = asset + } + if shared := byID[sharedID]; shared == nil || len(shared.TaskIDs) != 1 || shared.TaskIDs[0] != second.ID { + t.Fatalf("deleted task should be detached from shared asset: %+v", shared) + } + if protected := byID[protectedID]; protected == nil || len(protected.TaskIDs) != 0 { + t.Fatalf("other task's anchor should preserve the asset after detaching ownership: %+v", protected) + } +} + +func TestTaskRelationsAndLLMFailoverChain(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileIDs := make([]int64, 0, 3) + for i := 0; i < 3; i++ { + id, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("task-chain-%d-%d", suffix, i), Format: "openai", + Model: fmt.Sprintf("model-%d", i), APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + profileIDs = append(profileIDs, id) + } + + sourceA, err := d.CreateTask("source A", "goal A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + sourceB, err := d.CreateTask("source B", "goal B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + child, err := d.CreateTaskWithOptions("child", "new goal", TaskCreateOptions{ + SourceTaskIDs: []int64{sourceA.ID, sourceB.ID}, + LLMProfileIDs: profileIDs, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(child.ID) + _ = d.DeleteTask(sourceA.ID) + _ = d.DeleteTask(sourceB.ID) + for _, id := range profileIDs { + _ = d.DeleteProfile(id) + } + }) + + got, err := d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceA.ID, sourceB.ID}) { + t.Fatalf("unexpected direct sources: %v", got.SourceTaskIDs) + } + if fmt.Sprint(got.LLMProfileIDs) != fmt.Sprint(profileIDs) || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("unexpected initial chain: %+v", got) + } + initialRevision := got.LLMChainRevision + if err := d.ReplaceTaskLLMProfiles(child.ID, profileIDs, profileIDs[0]); err != nil { + t.Fatal(err) + } + staleRevision, err := d.MarkTaskLLMProfileQuotaExhaustedAtRevision(child.ID, profileIDs[0], initialRevision, "late quota from prior generation") + if err != nil { + t.Fatal(err) + } + if !staleRevision.Stale || staleRevision.Advanced || staleRevision.ChainExhausted { + t.Fatalf("old chain revision changed replacement chain: %+v", staleRevision) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("replacement revision/cursor not preserved: %+v", got) + } + replacementRevision := got.LLMChainRevision + + transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota") + if err != nil { + t.Fatal(err) + } + if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] { + t.Fatalf("unexpected first transition: %+v", transition) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.LLMChainRevision != replacementRevision+1 { + t.Fatalf("automatic cursor advance did not increment revision: before=%d after=%d", replacementRevision, got.LLMChainRevision) + } + late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate") + if err != nil { + t.Fatal(err) + } + if late.Advanced || late.NextProfileID == nil || *late.NextProfileID != profileIDs[1] { + t.Fatalf("late failure advanced past current profile: %+v", late) + } + if _, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "quota_exceeded"); err != nil { + t.Fatal(err) + } + last, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "余额不足") + if err != nil { + t.Fatal(err) + } + if !last.Advanced || !last.ChainExhausted || last.NextProfileID != nil { + t.Fatalf("unexpected exhausted transition: %+v", last) + } + duplicateLast, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "late final duplicate") + if err != nil { + t.Fatal(err) + } + if duplicateLast.Advanced || !duplicateLast.ChainExhausted || duplicateLast.NextProfileID != nil { + t.Fatalf("duplicate final failure must be idempotent: %+v", duplicateLast) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("chain exhaustion was not persisted: %+v", got) + } + + // Two requests can observe the same last active profile. Only the transaction + // that clears the cursor may report an advance; the late one is idempotent. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2]}, profileIDs[2]); err != nil { + t.Fatal(err) + } + transitions := make(chan TaskLLMTransition, 2) + errs := make(chan error, 2) + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + transition, markErr := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[2], "concurrent final quota") + if markErr != nil { + errs <- markErr + return + } + transitions <- transition + }() + } + wg.Wait() + close(errs) + close(transitions) + for markErr := range errs { + t.Fatal(markErr) + } + advanced := 0 + for transition := range transitions { + if transition.Advanced { + advanced++ + } + if !transition.ChainExhausted { + t.Fatalf("concurrent final transition must report exhausted chain: %+v", transition) + } + } + if advanced != 1 { + t.Fatalf("last profile advanced %d times, want exactly once", advanced) + } + + // Starting manually from the middle consumes only candidates after that cursor. + // Earlier ready profiles must not be revived by hydration or unrelated deletes. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil { + t.Fatal(err) + } + if step, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "middle quota"); err != nil || step.NextProfileID == nil || *step.NextProfileID != profileIDs[1] { + t.Fatalf("middle cursor did not advance to its successor: transition=%+v err=%v", step, err) + } + if end, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[1], "tail quota"); err != nil || !end.ChainExhausted { + t.Fatalf("tail did not exhaust manual chain: transition=%+v err=%v", end, err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("hydration revived a profile before the manual cursor: %+v", got) + } + unrelatedID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("task-chain-unrelated-%d", suffix), Format: "openai", Model: "other", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + if err := d.DeleteProfile(unrelatedID); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || got.LLMFailoverState != "chain_exhausted" { + t.Fatalf("deleting an unrelated profile revived an exhausted chain: %+v", got) + } + + // A call that finishes after its profile was removed from the chain is stale: + // preserve the new cursor and surface the original provider error upstream. + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[1]}, profileIDs[2]); err != nil { + t.Fatal(err) + } + stale, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "obsolete in-flight quota") + if err != nil { + t.Fatalf("removed in-flight profile returned an internal error: %v", err) + } + if stale.Advanced || stale.ChainExhausted || stale.NextProfileID != nil { + t.Fatalf("removed in-flight profile changed the replacement chain: %+v", stale) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[2] { + t.Fatalf("stale failure changed active replacement profile: %+v", got) + } + + if err := d.ReplaceTaskLLMProfiles(child.ID, []int64{profileIDs[2], profileIDs[0], profileIDs[1]}, profileIDs[0]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] || got.LLMFailoverState != "ready" || got.LLMFailoverReason != "" { + t.Fatalf("chain reset did not clear failure state: %+v", got) + } + + if err := d.DeleteProfile(profileIDs[0]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[1] { + t.Fatalf("deleting active profile did not select next ready entry: %+v", got) + } + if err := d.DeleteProfile(profileIDs[1]); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMFailoverState != "default" { + t.Fatalf("deleting active chain tail must fall back instead of wrapping: %+v", got) + } + + if err := d.DeleteTask(sourceA.ID); err != nil { + t.Fatal(err) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceB.ID}) { + t.Fatalf("source delete did not cascade only its relation: %v", got.SourceTaskIDs) + } +} + +func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + source, err := d.CreateTask("source", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(source.ID) }) + if _, err := d.CreateTaskWithOptions("bad", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID, source.ID}}); err == nil { + t.Fatal("duplicate source task ids should be rejected") + } + + task, err := d.CreateTask("terminal", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + if err := d.SetStatus(task.ID, "done"); err != nil { + t.Fatal(err) + } + if err := d.ReplaceTaskLLMProfiles(task.ID, nil, 0); err == nil { + t.Fatal("terminal task LLM edit should be rejected") + } +} + +func TestCreateTaskWithCompanyScopes(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "") + if err != nil { + t.Fatal(err) + } + companyB, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company B %d", suffix), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB) + }) + + task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{ + CompanyIDs: []int64{companyA, companyB}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) { + t.Fatalf("creation result company IDs=%v", task.CompanyIDs) + } + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) { + t.Fatalf("hydrated company IDs=%v", got.CompanyIDs) + } + var scopeCount int + if err := d.QueryRow(`SELECT count(*) FROM task_scope WHERE task_id=$1 AND kind='company'`, task.ID).Scan(&scopeCount); err != nil { + t.Fatal(err) + } + if scopeCount != 2 { + t.Fatalf("company task scope rows=%d want 2", scopeCount) + } + + duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{ + CompanyIDs: []int64{companyA, companyA, companyB, companyA}, + }) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(duplicateTask.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) { + t.Fatalf("company IDs were not normalized: %v", duplicateTask.CompanyIDs) + } + if err := d.DeleteTask(duplicateTask.ID); err != nil { + t.Fatal(err) + } + + badDescription := fmt.Sprintf("invalid-company-%d", suffix) + if _, err := d.CreateTaskWithOptions(badDescription, "rollback", TaskCreateOptions{CompanyIDs: []int64{companyA, 1 << 62}}); !errors.Is(err, ErrTaskCompanyNotFound) { + t.Fatalf("missing company error=%v, want %v", err, ErrTaskCompanyNotFound) + } + var leaked int + if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, badDescription).Scan(&leaked); err != nil { + t.Fatal(err) + } + if leaked != 0 { + t.Fatalf("failed company association leaked %d exploration rows", leaked) + } +} diff --git a/server/chatupload.go b/server/chatupload.go index 7ac55e1..82ae080 100644 --- a/server/chatupload.go +++ b/server/chatupload.go @@ -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()) diff --git a/server/engine.go b/server/engine.go index 699b204..551aa0a 100644 --- a/server/engine.go +++ b/server/engine.go @@ -71,8 +71,18 @@ func preview(s string, n int) string { const ( modelErrorRetries = 2 // model_error 收场后额外重试的次数 modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避 + workControlWaitTimeout = 30 * time.Second ) +var errWorkControlConflict = errors.New("work control conflict") + +// 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,25 +96,40 @@ 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) + 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. + // 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.CancelFunc + 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.CancelFunc + // 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 @@ -121,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 error + action string // user action: pause | cancel } // nextPlannerRound returns the next planner round number for a task (1-based). @@ -136,18 +175,128 @@ 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). -func (e *Engine) Pause(taskID string) { +func (e *Engine) Pause(taskID string, cause error) { e.paused.Store(taskID, true) - e.cancelExec(taskID) + e.cancelExec(taskID, cause) +} + +// BeginDelete installs an execution barrier before task data/files are removed. +// The temporary pause is not a user pause. The server serializes this transition +// with lifecycle admission and tells AbortDelete whether the persisted task is +// paused/queued if 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 + } + e.paused.Store(taskID, true) + e.deleteMu.Unlock() + e.cancelExec(taskID, agent.AbortTaskDeleted) + return true +} + +func (e *Engine) AbortDelete(taskID string, keepPaused bool) { + e.deleteMu.Lock() + if !e.IsDeleting(taskID) { + e.deleteMu.Unlock() + return + } + e.deleting.Delete(taskID) + if !keepPaused { + e.paused.Delete(taskID) + } + e.deleteMu.Unlock() + if !keepPaused && 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.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. -func (e *Engine) cancelExec(taskID string) { +func (e *Engine) cancelExec(taskID string, cause error) { e.execMu.Lock() if cancel := e.execCancel[taskID]; cancel != nil { - cancel() + cancel(cause) } e.execMu.Unlock() } @@ -155,6 +304,17 @@ func (e *Engine) cancelExec(taskID string) { // Resume un-pauses a task and nudges a fresh planning round. The next exec under // it gets a fresh (uncancelled) context. func (e *Engine) Resume(t *Task) { + // BeginDelete owns the pause barrier once deletion starts. A concurrent + // resume must never clear it and let a planner/worker re-enter while cleanup + // is waiting for task operations to drain. + if t == nil { + return + } + e.deleteMu.RLock() + defer e.deleteMu.RUnlock() + if e.IsDeleting(t.ID) { + return + } e.paused.Delete(t.ID) t.Notify() } @@ -166,14 +326,14 @@ func (e *Engine) execContextFor(parent context.Context, taskID string) context.C defer e.execMu.Unlock() if e.IsPaused(taskID) { // never hand out a live context while paused (guards the claim→Execute race) - c, cancel := context.WithCancel(parent) - cancel() + c, cancel := context.WithCancelCause(parent) + cancel(agent.AbortPausedRaceGuard) return c } if c := e.execCtx[taskID]; c != nil && c.Err() == nil { return c } - c, cancel := context.WithCancel(parent) + c, cancel := context.WithCancelCause(parent) e.execCtx[taskID] = c e.execCancel[taskID] = cancel return c @@ -200,32 +360,131 @@ 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.CancelFunc{}, execCtx: map[string]context.Context{}, - workCancel: map[int64]context.CancelFunc{}, steerBox: map[int64][]string{}} + execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{}, + 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.CancelFunc) { +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 error, 1)} 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(error)) { e.workMu.Lock() - if c := e.workCancel[intentID]; c != nil { - delete(e.workCancel, intentID) - c() // 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(error) {} + } + return action, func(err error) { run.done <- err } +} + +// 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(ctx context.Context, intentID int64, action string) error { + if action != "pause" && action != "cancel" { + return fmt.Errorf("unsupported work action %q", action) + } + if ctx == nil { + ctx = context.Background() + } + e.workMu.Lock() + run := e.work[intentID] + if run == nil { + e.workMu.Unlock() + return fmt.Errorf("%w: 意图 %d 当前没有运行中的 work（可能已结束或未被领取）", errWorkControlConflict, intentID) + } + if run.action != "" { + e.workMu.Unlock() + return fmt.Errorf("%w: 意图 %d 正在执行 %s 操作", errWorkControlConflict, 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() + + timer := time.NewTimer(workControlWaitTimeout) + defer timer.Stop() + select { + case err := <-done: + return err + case <-ctx.Done(): + e.releaseWorkControl(intentID, run, action) + return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, ctx.Err()) + case <-timer.C: + e.releaseWorkControl(intentID, run, action) + return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, context.DeadlineExceeded) + } +} + +// releaseWorkControl drops only this caller's reservation after its wait is +// cancelled. The work context stays cancelled; runWorkerStep recognizes the +// named cancellation cause and settles the intent into the recoverable paused +// state even if the HTTP caller has gone away. +func (e *Engine) releaseWorkControl(intentID int64, run *workExecution, action string) { + e.workMu.Lock() + if current := e.work[intentID]; current == run && current.action == action { + current.action = "" + } + e.workMu.Unlock() +} + +func transitionIntentState(store *db.ExplorationStore, intentID int64, expected, state string) error { + changed, err := store.CompareAndSetIntentState(intentID, expected, state) + if err != nil { + return err + } + if !changed { + return fmt.Errorf("%w: 意图 %d 不再是 %s 状态", db.ErrIntentStateConflict, intentID, expected) + } + return nil } // SteerWork queues a mid-run course-correction for the work running intentID (the @@ -236,7 +495,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) @@ -301,12 +560,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() + run.cancel(agent.AbortKilledByPlanner) return nil } @@ -315,7 +574,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 @@ -340,7 +599,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() { @@ -348,6 +607,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. @@ -390,6 +650,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 @@ -397,7 +666,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 } } @@ -410,6 +680,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() { @@ -422,15 +700,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 首轮规划。带种子意图的任务:种子已 @@ -491,10 +778,13 @@ 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). - if isTerminalStatus(t.Status) { + if isTerminalStatus(t.lifecycleSnapshot().Status) { return } // 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再 @@ -502,10 +792,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'). @@ -514,10 +811,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) @@ -530,7 +827,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { // 任务已判完成 → 立刻取消在跑的 worker：它们手头的意图跑出来也没意义了。 // 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支 // 归为 stopped(而非 blocked)。 - e.cancelExec(t.ID) + e.cancelExec(t.ID, agent.AbortGoalMet) default: log.Printf("[planner] task %s 规划完成", t.ID) } @@ -572,6 +869,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 @@ -584,105 +884,191 @@ 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.WithCancel(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 + } +} + +// 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) { + if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil { + log.Printf("[worker %s] task %s 意图 #%d 领取后回退失败: %v", name, t.ID, intent.ID, err) } - // 任务已判完成(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 + 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 // 退避期间被取消（终止/暂停）→ 交给下方分支处理 } - // 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" + 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. + workCause := context.Cause(workCtx) + killed := workCtx.Err() != nil && ectx.Err() == nil + action, completeWork := e.detachWork(intent.ID) + // A caller may stop waiting and release its in-memory reservation before the + // agent honors cancellation. The named context cause remains authoritative and + // still settles the stopped run into a recoverable state. + if action == "" { 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) + case errors.Is(workCause, agent.AbortWorkPausedByUser): + action = "pause" + case errors.Is(workCause, agent.AbortWorkCancelledByUser): + action = "cancel" } - _ = t.Store.SetIntentState(intent.ID, state) - log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote) + } + var controlErr error + defer func() { completeWork(controlErr) }() + if action == "pause" { + controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused") + if controlErr != nil { + log.Printf("[worker %s] task %s 意图 #%d 暂停状态落库失败: %v", name, t.ID, intent.ID, controlErr) + return true + } + log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID) e.touch(t.ID) - t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id) + return true + } + if action == "cancel" { + // Park the stopped run in paused before handing cleanup to the API. If the + // request disconnects after cancellation, the intent remains recoverable and + // a later cancel can finish cleanup instead of leaving a phantom running row. + controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused") + if controlErr != nil { + log.Printf("[worker %s] task %s 意图 #%d 取消栅栏落库失败: %v", name, t.ID, intent.ID, controlErr) + return true + } + 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 on resume — the worker will resume the prior LLM + // conversation from its transcript instead of restarting from scratch. + if ectx.Err() != nil && taskExecutionPaused(context.Cause(ectx)) { + if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil { + log.Printf("[worker %s] task %s 意图 #%d 任务暂停回退失败: %v", name, t.ID, intent.ID, err) + } + return true + } + // 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾), + // 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。 + if ectx.Err() != nil && e.isSettling(t.ID) { + if err := transitionIntentState(t.Store, intent.ID, "running", "exhausted"); err != nil { + log.Printf("[worker %s] task %s 意图 #%d 超时收尾状态落库失败: %v", name, t.ID, intent.ID, err) + } + 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)) { + if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil { + log.Printf("[worker %s] task %s 意图 #%d 终态停止落库失败: %v", name, t.ID, intent.ID, err) + } + 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 { + if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil { + log.Printf("[worker %s] task %s 意图 #%d planner 停止落库失败: %v", name, t.ID, intent.ID, err) + } + 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 { + if stateErr := transitionIntentState(t.Store, intent.ID, "running", state); stateErr != nil { + log.Printf("[worker %s] task %s 意图 #%d 终态 %s 落库失败: %v", name, t.ID, intent.ID, state, stateErr) + } + } + 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 taskExecutionPaused(cause error) bool { + var abort *agent.AbortCause + if !errors.As(cause, &abort) { + return false + } + switch abort.Code { + case "paused_by_user", "paused_by_orchestrator", "paused_on_reload", "paused_race_guard", + "queued_for_admission", "llm_unavailable_queued", "task_deleted": + return true + default: + return false } } diff --git a/server/engine_cancelcause_test.go b/server/engine_cancelcause_test.go new file mode 100644 index 0000000..d1f4f7a --- /dev/null +++ b/server/engine_cancelcause_test.go @@ -0,0 +1,217 @@ +package server + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" +) + +func TestPauseCancelsWithNamedCause(t *testing.T) { + e := NewEngine(nil) + 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("code=%q ok=%v, want paused_by_user", code, ok) + } +} + +func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) { + e := NewEngine(nil) + execCtx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(execCtx) + e.registerWork(42, workCancel) + if err := e.KillWork(42); err != nil { + t.Fatal(err) + } + if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "killed_by_planner" { + t.Fatalf("code=%q ok=%v, want killed_by_planner", code, ok) + } + if execCtx.Err() != nil { + t.Fatal("kill_work cancelled the entire task context") + } + _, complete := e.detachWork(42) + complete(nil) +} + +func TestControlWorkCarriesUserActionCause(t *testing.T) { + for _, tc := range []struct { + action string + code string + }{ + {action: "pause", code: "work_paused_by_user"}, + {action: "cancel", code: "work_cancelled_by_user"}, + } { + t.Run(tc.action, func(t *testing.T) { + e := NewEngine(nil) + ctx, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + done := make(chan error, 1) + go func() { done <- e.ControlWork(context.Background(), 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(nil) + if err := <-done; err != nil { + t.Fatal(err) + } + }) + } +} + +func TestControlWorkReturnsSettlementError(t *testing.T) { + e := NewEngine(nil) + ctx, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + done := make(chan error, 1) + go func() { done <- e.ControlWork(context.Background(), 42, "pause") }() + <-ctx.Done() + _, complete := e.detachWork(42) + want := errors.New("pause persistence failed") + complete(want) + if err := <-done; !errors.Is(err, want) { + t.Fatalf("ControlWork error=%v, want %v", err, want) + } +} + +func TestControlWorkWaitHonorsContextAndReleasesReservation(t *testing.T) { + e := NewEngine(nil) + _, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + waitCtx, stop := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer stop() + if err := e.ControlWork(waitCtx, 42, "pause"); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("ControlWork error=%v, want deadline exceeded", err) + } + + e.workMu.Lock() + run := e.work[42] + if run == nil || run.action != "" { + t.Fatalf("timed-out control left reservation: %+v", run) + } + e.workMu.Unlock() + _, complete := e.detachWork(42) + complete(nil) +} + +func TestControlWorkRejectsConcurrentController(t *testing.T) { + e := NewEngine(nil) + ctx, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + first := make(chan error, 1) + go func() { first <- e.ControlWork(context.Background(), 42, "pause") }() + <-ctx.Done() + if err := e.ControlWork(context.Background(), 42, "cancel"); !errors.Is(err, errWorkControlConflict) { + t.Fatalf("second control error=%v, want conflict", err) + } + action, complete := e.detachWork(42) + if action != "pause" { + t.Fatalf("winning action=%q, want pause", action) + } + complete(nil) + if err := <-first; err != nil { + t.Fatal(err) + } +} + +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("paused task received a live execution context") + } + if code, _, _, _ := agent.AbortReason(ctx); code != "paused_race_guard" { + t.Fatalf("code=%q, want paused_race_guard", code) + } +} + +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 did not cancel old context") + } + e.paused.Delete("t1") + if fresh := e.execContextFor(context.Background(), "t1"); fresh.Err() != nil { + t.Fatal("resume did not create a fresh execution context") + } +} + +func TestTaskPauseCauseSurvivesImmediateResume(t *testing.T) { + e := NewEngine(nil) + task := &Task{ID: "t1", notify: make(chan struct{}, 1)} + old := e.execContextFor(context.Background(), task.ID) + e.Pause(task.ID, agent.AbortPausedByUser) + e.Resume(task) + if e.IsPaused(task.ID) { + t.Fatal("resume did not clear the task pause flag") + } + if old.Err() == nil || !taskExecutionPaused(context.Cause(old)) { + t.Fatalf("old execution cause=%v was not retained as a task pause", context.Cause(old)) + } + if taskExecutionPaused(agent.AbortGoalMet) { + t.Fatal("goal completion was misclassified as a task pause") + } +} + +func TestDeleteCancellationReturnsRunningIntentToFrontier(t *testing.T) { + if !taskExecutionPaused(agent.AbortTaskDeleted) { + t.Fatal("delete cancellation must settle a running intent as open so AbortDelete can resume it") + } + if taskExecutionPaused(agent.AbortGoalMet) { + t.Fatal("terminal task cancellation must not reopen completed work") + } +} + +func TestResumeDoesNotClearDeleteBarrier(t *testing.T) { + e := NewEngine(nil) + task := &Task{ID: "t1", notify: make(chan struct{}, 1)} + e.paused.Store(task.ID, true) + e.deleting.Store(task.ID, true) + e.Resume(task) + if !e.IsPaused(task.ID) { + t.Fatal("resume cleared the pause barrier owned by task deletion") + } +} + +func TestBootstrapAdmissionClearsQueuedPauseBarrier(t *testing.T) { + e := NewEngine(nil) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + task := &Task{ID: "t1", notify: make(chan struct{}, 1)} + s := &Server{engine: e, ctx: ctx} + e.Pause(task.ID, agent.AbortPausedByUser) + + s.startAdmittedTask(task, "bootstrap") + if e.IsPaused(task.ID) { + t.Fatal("bootstrap admission left the queued task behind the pause barrier") + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 14fa3fe..8d57293 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -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)) @@ -58,13 +71,26 @@ func (e *Engine) taskDeadline(t *Task) int64 { if v, ok := e.deadline.Load(t.ID); ok { return v.(int64) } - if t.DeadlineAt > 0 { - e.deadline.Store(t.ID, t.DeadlineAt) - return t.DeadlineAt + deadlineAt := t.lifecycleSnapshot().DeadlineAt + if deadlineAt > 0 { + e.deadline.Store(t.ID, deadlineAt) + return deadlineAt } return 0 } +// resetTimeoutRevival clears only the per-run timeout state after PostgreSQL has +// atomically committed timeout -> running and reset first_run_at/deadline_at. The +// configured TimeoutSeconds remains on Task, so the next real Planner/Worker run +// stamps a fresh full budget. coordStarted is reset because the coordinator that +// produced the timeout has already completed (or is in its final return path). +func (e *Engine) resetTimeoutRevival(taskID string) { + e.settling.Delete(taskID) + e.deadline.Delete(taskID) + e.stamped.Delete(taskID) + e.coordStarted.Delete(taskID) +} + // stampFirstRun records first_run_at + deadline_at on the FIRST real run (LLM ready) // of a timeout task, once per process. No-op when the task has no timeout. func (e *Engine) stampFirstRun(t *Task) { @@ -106,10 +132,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 @@ -159,7 +193,7 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { for e.inflightCount(t.ID) > 0 { if time.Now().After(hardStop) { log.Printf("[deadline] task %s drain 超时(%s),硬取消在跑 run", t.ID, settleDrainGrace) - e.cancelExec(t.ID) + e.cancelExec(t.ID, agent.AbortSettleDrainTimeout) _ = sleepCtx(ctx, 3*time.Second) // 给 worker 分支一点时间落库/归类 break } @@ -170,6 +204,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 +230,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 +241,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 { diff --git a/server/goals.go b/server/goals.go index 09e23b1..dbc497d 100644 --- a/server/goals.go +++ b/server/goals.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "log" + "sort" "strconv" "strings" @@ -25,25 +26,324 @@ 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) } - go func() { - s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", - Summary: "第 0 轮目标拆解"}) - goals := s.createGoals(s.ctx, t, func(r db.Activity) { - s.engine.emitActivity(t, r) - }) - for _, g := range goals { - summary := g.Text - if g.VulnClass != "" { - summary = fmt.Sprintf("[%s] %s", g.VulnClass, g.Text) - } - s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary}) + s.engine.decInflight(t.ID) + if _, err := s.admitTask(t, "bootstrap"); err != nil { + log.Printf("[concurrency] task %s 启动失败: %v", t.ID, err) + } +} + +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(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 != "" { + summary = fmt.Sprintf("[%s] %s", g.VulnClass, g.Text) } - s.engine.Run(s.ctx, t) - }() + 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) +} + +func (s *Server) occupiesConcurrencySlot(t *Task) bool { + if t == nil { + return false + } + lifecycle := t.lifecycleSnapshot() + if lifecycle.Queued || lifecycle.Paused || isTerminalStatus(lifecycle.Status) { + return false + } + // Deletion temporarily pauses the Engine but has not committed yet. Preserve + // the task's slot until PostgreSQL deletion succeeds; otherwise FIFO promotion + // during the drain window could over-admit if deletion later aborts and the + // persisted running task is restored. + if s.engine.IsDeleting(t.ID) { + return true + } + return !s.engine.IsPaused(t.ID) && (s.engine.ReadyFor(t) || s.engine.ActiveLLMCalls(t.ID) > 0) +} + +func (s *Server) runningTaskCount(excludeID string) int { + count := 0 + for _, task := range s.m.List() { + if task.ID != excludeID && s.occupiesConcurrencySlot(task) { + count++ + } + } + return count +} + +// admitTask is the single admission path for new, resumed, rerun and follow-up +// work. It atomically either starts the task or appends it to the persistent FIFO +// queue. mode is bootstrap for a freshly-created task and resume otherwise. +func (s *Server) admitTask(t *Task, mode string) (queued bool, err error) { + return s.admitTaskWhen(t, mode, false) +} + +// admitPausedTask is the task-control resume path. The paused precondition is +// checked under the same scheduler lock as admission so a concurrent pause, +// dequeue or FIFO promotion cannot leave database and Engine state divergent. +func (s *Server) admitPausedTask(t *Task) (queued bool, err error) { + return s.admitTaskWhen(t, "resume", true) +} + +func (s *Server) admitTaskWhen(t *Task, mode string, requirePaused bool) (queued bool, err error) { + if t == nil { + return false, fmt.Errorf("task not found") + } + if mode != "bootstrap" { + mode = "resume" + } + s.concMu.Lock() + defer s.concMu.Unlock() + // Delete installs its barrier under concMu as well. Re-resolve after acquiring + // the lock so a request that captured a task pointer before successful deletion + // cannot revive that stale handle after StopTask clears its Engine maps. + current, exists := s.m.Task(t.ID) + if !exists || current != t || s.engine.IsDeleting(t.ID) { + return false, fmt.Errorf("task is being deleted") + } + if !s.engine.beginTaskOperation(t.ID) { + return false, fmt.Errorf("task is being deleted") + } + defer s.engine.decInflight(t.ID) + lifecycle := t.lifecycleSnapshot() + if requirePaused { + if isTerminalStatus(lifecycle.Status) { + return false, fmt.Errorf("终态任务不能执行继续") + } + if !lifecycle.Paused { + return false, fmt.Errorf("仅已暂停的任务可以继续") + } + mode = s.resumeAdmissionMode(t) + } + wasTerminal := isTerminalStatus(lifecycle.Status) + wasPaused := lifecycle.Paused || s.engine.IsPaused(t.ID) + wasQueued := lifecycle.Queued + engineWasPaused := s.engine.IsPaused(t.ID) + enabled, limit := s.m.ConcurrencyLimit() + ready := s.engine.ReadyFor(t) + + // Work added to a task that is already admitted only needs a wake-up. This + // matters when the configured limit was lowered below the current running + // count: an existing task must not suddenly mark itself queued while its + // planner/workers are still live. Still compare-and-commit the persisted + // status: a concurrent terminal transition must win instead of being silently + // reported as a successful follow-up admission. + if !wasTerminal && !wasPaused && !wasQueued && s.engine.Started(t.ID) && (!enabled || ready) { + if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, lifecycle.Status, false, "resume", false); err != nil { + return false, err + } + s.startAdmittedTask(t, "resume") + return false, nil + } + + // Preserve the original first-run mode across repeated admissions. Legacy + // queued rows have an empty queue_mode, so infer bootstrap from their graph. + if wasQueued { + switch lifecycle.QueueMode { + case "bootstrap": + mode = "bootstrap" + case "": + mode = s.resumeAdmissionMode(t) + } + } + + readyBacklog := enabled && s.hasReadyQueuedTask(t.ID) + atCapacity := enabled && s.runningTaskCount(t.ID) >= limit + shouldQueue := enabled && (!ready || readyBacklog || atCapacity) + + // Install the execution barrier before reviving a terminal/paused task. Without + // this ordering, its already-running worker loops can claim the newly-opened + // intent in the gap between status=running and queued=true. + if shouldQueue || wasTerminal || wasPaused || wasQueued { + s.engine.Pause(t.ID, agent.Causef("queued_for_admission", "任务等待运行准入", + "任务正在等待并发队列或准入状态提交，本次执行已停止；只有获得运行槽后才会重新领取意图")) + } + + status := lifecycle.Status + if wasTerminal { + status = "running" + } + if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, status, shouldQueue, mode, wasQueued); err != nil { + if !engineWasPaused && !wasPaused && !wasQueued { + s.engine.Resume(t) + } + return false, err + } + if lifecycle.Status == "timeout" && status == "running" { + // ApplyTaskAdmission reset this timed-out run's persisted clock. Clear the + // matching Engine gates before either parking it in FIFO or starting work; + // otherwise the old settling flag would make every worker skip forever. + s.engine.resetTimeoutRevival(t.ID) + } + if shouldQueue { + if !wasQueued { + summary := fmt.Sprintf("已排队：达到并发上限 %d，等待空位后自动开始", limit) + switch { + case !ready: + summary = "已排队：当前没有可运行的 LLM 配置，配置恢复后自动开始" + case readyBacklog: + summary = "已排队：已有更早的任务等待运行，将按 FIFO 顺序自动开始" + } + s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: summary}) + } + return true, nil + } + s.startAdmittedTask(t, mode) + return false, nil +} + +func (s *Server) hasReadyQueuedTask(excludeID string) bool { + for _, task := range s.m.List() { + lifecycle := task.lifecycleSnapshot() + if task.ID != excludeID && lifecycle.Queued && !lifecycle.Paused && !isTerminalStatus(lifecycle.Status) && s.engine.ReadyFor(task) { + return true + } + } + return false +} + +func (s *Server) startAdmittedTask(t *Task, mode string) { + if mode == "bootstrap" { + if !s.engine.beginTaskOperation(t.ID) { + return + } + // A first-run task may have kept the Engine pause barrier while waiting + // in the concurrency queue. Clear it only after operation admission, or + // startTaskEngine's first execContextFor call would return a cancelled + // context and silently strand the dequeued task. + if s.engine.IsPaused(t.ID) { + s.engine.Resume(t) + } + go func() { + defer s.engine.decInflight(t.ID) + s.startTaskEngine(t) + }() + return + } + s.engine.Run(s.ctx, t) + s.engine.Resume(t) + // Run returns early for an already-started task. Explicitly ensure a timeout + // coordinator exists after resume; resetTimeoutRevival cleared the completed + // coordinator's marker and the next real call will stamp a fresh deadline. + s.engine.startDeadlineCoordinator(s.ctx, t) +} + +func (s *Server) reconcileConcurrency() { + s.concMu.Lock() + defer s.concMu.Unlock() + enabled, limit := s.m.ConcurrencyLimit() + + // A task whose provider chain becomes unavailable cannot do useful work and + // must not reserve a limited running slot forever. Park it persistently so a + // later profile edit/recovery can re-enter through the same FIFO path. + if enabled { + for _, task := range s.m.List() { + lifecycle := task.lifecycleSnapshot() + if lifecycle.Queued || lifecycle.Paused || s.engine.IsPaused(task.ID) || s.engine.IsDeleting(task.ID) || + isTerminalStatus(lifecycle.Status) || s.engine.ReadyFor(task) || s.engine.ActiveLLMCalls(task.ID) > 0 { + continue + } + mode := s.resumeAdmissionMode(task) + s.engine.Pause(task.ID, agent.Causef("llm_unavailable_queued", "LLM 不可用，任务进入等待队列", + "任务当前无法解析可运行的 Planner/Worker LLM，已释放并发槽；配置恢复后按队列顺序继续")) + if err := s.m.EnqueueTask(task.ID, mode); err != nil { + s.engine.Resume(task) + log.Printf("[concurrency] task %s 因 LLM 不可用入队失败: %v", task.ID, err) + continue + } + s.engine.emitActivity(task, db.Activity{Worker: "system", Kind: "text", + Summary: "已排队：当前没有可运行的 LLM 配置，配置恢复后自动开始"}) + } + } + + type queuedTask struct { + task *Task + lifecycle taskLifecycleState + } + queued := []queuedTask{} + for _, task := range s.m.List() { + lifecycle := task.lifecycleSnapshot() + if lifecycle.Queued && !isTerminalStatus(lifecycle.Status) { + queued = append(queued, queuedTask{task: task, lifecycle: lifecycle}) + } + } + sort.SliceStable(queued, func(i, j int) bool { + left, right := queued[i].lifecycle.QueuedAt, queued[j].lifecycle.QueuedAt + if left == 0 { + left = queued[i].task.CreatedAt * int64(1e9) + } + if right == 0 { + right = queued[j].task.CreatedAt * int64(1e9) + } + if left != right { + return left < right + } + // Old rows may not have queued_at and task creation timestamps are only + // kept to second precision in memory. Task ids are monotonic, providing a + // deterministic oldest-first fallback for those ties. + leftID, leftErr := strconv.ParseInt(queued[i].task.ID, 10, 64) + rightID, rightErr := strconv.ParseInt(queued[j].task.ID, 10, 64) + if leftErr == nil && rightErr == nil { + return leftID < rightID + } + return queued[i].task.ID < queued[j].task.ID + }) + slots := len(queued) + if enabled { + slots = limit - s.runningTaskCount("") + } + for _, entry := range queued { + if slots <= 0 { + break + } + task := entry.task + lifecycle := task.lifecycleSnapshot() + // 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 lifecycle.Paused || (enabled && !s.engine.ReadyFor(task)) { + continue + } + mode := lifecycle.QueueMode + if mode != "bootstrap" && mode != "resume" { + mode = s.resumeAdmissionMode(task) + } + if mode != "bootstrap" { + mode = "resume" + } + if !s.engine.beginTaskOperation(task.ID) { + continue + } + if err := s.m.ApplyTaskAdmission(task.ID, lifecycle.Status, lifecycle.Status, false, mode, false); err != nil { + s.engine.decInflight(task.ID) + continue + } + s.startAdmittedTask(task, mode) + s.engine.decInflight(task.ID) + slots-- + } } // reviveTask 让一个已停下的任务重新跑起来:把终态(done/failed/timeout)拉回 running、 @@ -56,22 +356,8 @@ func (s *Server) reviveTask(t *Task) { if t == nil { return } - // 终态 → 拉回 running(SetTaskStatus 会清 completed_at 并同步内存 handle 的 Status, - // 使 planner/worker 循环的终态门放行)。 - if isTerminalStatus(t.Status) { - if err := s.m.SetTaskStatus(t.ID, "running"); err != nil { - log.Printf("[revive] task %s 置 running 失败: %v", t.ID, err) - } - } - // 确保引擎循环存活:已 started 只会 Notify;未 started(重启后未恢复的终态任务)则重起循环。 - s.engine.Run(s.ctx, t) - // 暂停中 → 解除暂停(清引擎内存态 + Notify),并持久化 paused=false(存活过重启)。 - if t.Paused || s.engine.IsPaused(t.ID) { - s.engine.Resume(t) - t.Paused = false - if err := s.m.SetTaskPaused(t.ID, false); err != nil { - log.Printf("[revive] task %s 解除暂停失败: %v", t.ID, err) - } + if _, err := s.admitTask(t, "resume"); err != nil { + log.Printf("[revive] task %s 恢复失败: %v", t.ID, err) } } @@ -83,24 +369,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 cfg, ok := s.goalsLLMConfig(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, cfg, 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 { @@ -119,19 +406,3 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity } return specs } - -// goalsLLMConfig resolves the LLM config for goal decomposition by the standard -// precedence: the goals agent's own binding → the task's pinned profile → the global -// active profile (same source the engine runs on). -func (s *Server) goalsLLMConfig(t *Task) (agent.Config, bool) { - var pin *int64 - if t != nil { - pin = t.LLMProfileID - } - if eff := s.effectiveProfileForAgent("goals", pin); eff != nil { - if cfg, ok := s.loadProfileConfig(*eff); ok { - return cfg, true - } - } - return s.loadLLMConfig() -} diff --git a/server/goals_test.go b/server/goals_test.go new file mode 100644 index 0000000..057d85b --- /dev/null +++ b/server/goals_test.go @@ -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) + } +} diff --git a/server/inheritance_api_test.go b/server/inheritance_api_test.go new file mode 100644 index 0000000..df9fe4f --- /dev/null +++ b/server/inheritance_api_test.go @@ -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 +} diff --git a/server/inheritance_dto_test.go b/server/inheritance_dto_test.go new file mode 100644 index 0000000..078d233 --- /dev/null +++ b/server/inheritance_dto_test.go @@ -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") + } +} diff --git a/server/manager.go b/server/manager.go index ef4ca11..46cf066 100644 --- a/server/manager.go +++ b/server/manager.go @@ -1,7 +1,9 @@ package server import ( + "database/sql" "encoding/json" + "errors" "fmt" "log" "os" @@ -12,48 +14,30 @@ 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 { -<<<<<<< Updated upstream ID string `json:"id"` ExpID int64 `json:"exploration_id"` + Name string `json:"name"` // 可选任务名称;空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` + CategoryName string `json:"category_name,omitempty"` 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"` - 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 为终态；空/其它则由运行态推导) - // 任务级超时(见 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{} -======= - ID string `json:"id"` - ExpID int64 `json:"exploration_id"` - Name string `json:"name"` // 可选任务名称;空=未命名 - Description string `json:"description"` - Goal string `json:"goal"` - CreatedAt int64 `json:"created_at"` - CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成 - Paused bool `json:"paused"` - Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 + Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 // QueuedAt is an internal Unix-nanosecond ordering key. It is deliberately // finer than CreatedAt so several tasks enqueued in the same second retain // their real FIFO order. @@ -80,7 +64,6 @@ type Task struct { notify chan struct{} lifecycleMu sync.RWMutex llmMu sync.RWMutex ->>>>>>> Stashed changes // pendingTriggers accumulates the concrete changes (worker done / finding) that // fired planning rounds since the last one consumed them. The debounce coalesces @@ -89,21 +72,160 @@ type Task struct { pendingTriggers []agent.TriggerEvent } +// taskLifecycleState is an internally consistent view of the mutable task +// lifecycle and inherited-scope context. Callers must use lifecycleSnapshot and +// updateLifecycle instead of reading or writing the corresponding Task fields +// directly after the task has been published by Manager. +type taskLifecycleState struct { + Status string + Paused bool + Queued bool + QueuedAt int64 + QueueMode string + CompletedAt int64 + FirstRunAt int64 + DeadlineAt int64 + SourceTaskIDs []int64 + CompanyIDs []int64 + CategoryID *int64 + CategoryName string +} + +func (t *Task) lifecycleSnapshot() taskLifecycleState { + if t == nil { + return taskLifecycleState{} + } + t.lifecycleMu.RLock() + defer t.lifecycleMu.RUnlock() + return t.lifecycleSnapshotLocked() +} + +func (t *Task) lifecycleSnapshotLocked() taskLifecycleState { + return taskLifecycleState{ + Status: t.Status, + Paused: t.Paused, + Queued: t.Queued, + QueuedAt: t.QueuedAt, + QueueMode: t.QueueMode, + CompletedAt: t.CompletedAt, + FirstRunAt: t.FirstRunAt, + DeadlineAt: t.DeadlineAt, + SourceTaskIDs: append([]int64(nil), t.SourceTaskIDs...), + CompanyIDs: append([]int64(nil), t.CompanyIDs...), + CategoryID: cloneInt64Ptr(t.CategoryID), + CategoryName: t.CategoryName, + } +} + +func (t *Task) updateLifecycle(update func(*taskLifecycleState)) { + if t == nil || update == nil { + return + } + t.lifecycleMu.Lock() + state := t.lifecycleSnapshotLocked() + update(&state) + t.Status = state.Status + t.Paused = state.Paused + t.Queued = state.Queued + t.QueuedAt = state.QueuedAt + t.QueueMode = state.QueueMode + t.CompletedAt = state.CompletedAt + t.FirstRunAt = state.FirstRunAt + t.DeadlineAt = state.DeadlineAt + t.SourceTaskIDs = append(t.SourceTaskIDs[:0], state.SourceTaskIDs...) + t.CompanyIDs = append(t.CompanyIDs[:0], state.CompanyIDs...) + t.CategoryID = cloneInt64Ptr(state.CategoryID) + t.CategoryName = state.CategoryName + t.lifecycleMu.Unlock() +} + +type taskLLMState struct { + ProfileID *int64 + ProfileIDs []int64 + ActiveID *int64 + ChainRevision 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), + ChainRevision: t.LLMChainRevision, + FailoverState: t.LLMFailoverState, + FailoverReason: t.LLMFailoverReason, + } +} + +func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, revision int64, state, reason string) bool { + t.llmMu.Lock() + defer t.llmMu.Unlock() + if revision < t.LLMChainRevision { + return false + } + t.LLMProfileID = cloneInt64Ptr(profileID) + t.ActiveLLMProfileID = cloneInt64Ptr(activeID) + t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...) + t.LLMChainRevision = revision + t.LLMFailoverState = state + t.LLMFailoverReason = reason + return true +} + +// 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 - tasks map[string]*Task - active string - trafficOn bool // 流量捕获开关（默认关；settings.traffic_capture） - llmRecOn bool // LLM 录制开关（默认关；settings.llm_record） + companyMu sync.Mutex // serializes task/company-scope commits with live handle registration + // taskStateMu preserves commit order between PostgreSQL lifecycle writes and + // their in-memory mirrors. lifecycleMu makes snapshots race-free, but without + // this outer write lock an older request could commit first and publish last. + taskStateMu sync.Mutex + mu sync.RWMutex + tasks map[string]*Task + active string + trafficOn bool // 流量捕获开关（默认关；settings.traffic_capture） + llmRecOn bool // LLM 录制开关（默认关；settings.llm_record） // 联网搜索开关与来源（默认关；settings.web_search_*）。brave-free 需要 braveKey；tavily 需要 tavilyKey。 // webSearchProxy 是独立出口代理(http/https/socks5)，与记录流量的 MITM 代理无关。 webSearchOn bool @@ -123,12 +245,51 @@ const ( settingWebSearchProxy = "web_search_proxy" settingWorkers = "workers" settingLLMRecord = "llm_record" + // LLM 轮询(故障转移)。默认关闭——开启后走「全局激活配置」的 agent 在当前配置 + // 不可用(余额不足/key 失效/限流/服务异常)时自动切到下一个配置。 + // settingLLMPoolBindFallback 仅在轮询开启时有意义:默认关闭,即 agent/任务显式 + // 绑定了某个配置就只用它、失败即失败;开启后绑定的配置失败也会回落到轮询链。 + settingLLMPoolOn = "llm_pool_enabled" + settingLLMPoolBindFallback = "llm_pool_bind_fallback" + // 任务并发上限:开关 + 上限数。默认关闭;开启后默认上限 5(见 defaultConcurrencyLimit)。 + settingConcurrencyOn = "task_concurrency_enabled" + settingConcurrencyLimit = "task_concurrency_limit" // defaultWebSearchBackend is used when web search is on but no backend was picked. defaultWebSearchBackend = "ddgs" // defaultWorkers is the concurrent work-agent count when the setting is unset. defaultWorkers = 3 + // defaultConcurrencyLimit is the simultaneous-running-task cap when the feature + // is enabled but no explicit limit was saved. + defaultConcurrencyLimit = 5 ) +// ConcurrencyLimit returns whether the simultaneous-running-task cap is enabled and +// its limit (default 5 when enabled but unset). limit is always >=1 when enabled. +func (m *Manager) ConcurrencyLimit() (enabled bool, limit int) { + on, _, _ := m.pg.GetSetting(settingConcurrencyOn) + if strings.TrimSpace(on) != "true" { + return false, 0 + } + limit = defaultConcurrencyLimit + if v, ok, _ := m.pg.GetSetting(settingConcurrencyLimit); ok { + if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n >= 1 { + limit = n + } + } + return true, limit +} + +// SetConcurrency persists the running-task concurrency cap. limit<1 is clamped to 1. +func (m *Manager) SetConcurrency(enabled bool, limit int) error { + if limit < 1 { + limit = defaultConcurrencyLimit + } + if err := m.pg.SetSetting(settingConcurrencyLimit, strconv.Itoa(limit)); err != nil { + return err + } + return m.pg.SetSetting(settingConcurrencyOn, strconv.FormatBool(enabled)) +} + // Workers returns the configured concurrent work-agent count (default 3). Read // per-task at engine.Run, so a change applies to tasks started afterwards. func (m *Manager) Workers() int { @@ -181,6 +342,9 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { if err := pg.EnsureLLMRecordsTable(); err != nil { log.Printf("[llmrec] create table: %v", err) } + if err := pg.EnsureLLMUsageTable(); err != nil { + log.Printf("[llmusage] create table: %v", err) + } m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)} if proxyAddr != "" { tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr) @@ -270,6 +434,36 @@ func (m *Manager) SetLLMRecordEnabled(on bool) error { return nil } +// LLMPoolEnabled reports whether LLM failover ("轮询") is on (默认关； +// settings.llm_pool_enabled). Read when the provider chain is built (applyLLM), +// so a change requires a rebuild — putSettings does that. +func (m *Manager) LLMPoolEnabled() bool { + if m.pg == nil { + return false + } + return m.pg.GetBool(settingLLMPoolOn, false) +} + +// SetLLMPoolEnabled persists the failover toggle. Callers rebuild agents +// (applyLLM) afterwards so it takes effect. +func (m *Manager) SetLLMPoolEnabled(on bool) error { return m.pg.SetBool(settingLLMPoolOn, on) } + +// LLMPoolBindFallback reports whether an agent/task that is BOUND to a specific +// profile still falls back to the chain when that profile fails (默认关：绑定即 +// 独占，失败即失败). Only meaningful while LLMPoolEnabled. +func (m *Manager) LLMPoolBindFallback() bool { + if m.pg == nil { + return false + } + return m.pg.GetBool(settingLLMPoolBindFallback, false) +} + +// SetLLMPoolBindFallback persists the bound-profile fallback toggle. Callers +// rebuild agents (applyLLM) afterwards. +func (m *Manager) SetLLMPoolBindFallback(on bool) error { + return m.pg.SetBool(settingLLMPoolBindFallback, on) +} + // WebSearch returns the current web-search config: whether it is enabled, the // backend ("ddgs" | "brave-free" | "tavily"), the Brave API key, the Tavily API // key (each empty unless set), and the dedicated egress proxy (empty = direct). @@ -492,18 +686,27 @@ func unixOrZero(t *time.Time) int64 { return t.Unix() } +func unixNanoOrZero(t *time.Time) int64 { + if t == nil { + return 0 + } + return t.UnixNano() +} + func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task { return &Task{ ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID, -<<<<<<< Updated upstream - Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, -======= - Name: pt.Name, + Name: pt.Name, + CategoryID: cloneInt64Ptr(pt.CategoryID), CategoryName: pt.CategoryName, Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued, QueuedAt: unixNanoOrZero(pt.QueuedAt), QueueMode: pt.QueueMode, ->>>>>>> Stashed changes CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef, - LLMProfileID: pt.LLMProfileID, + LLMProfileID: pt.LLMProfileID, + LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID, + LLMChainRevision: pt.LLMChainRevision, + LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason, + SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...), + CompanyIDs: append([]int64(nil), pt.CompanyIDs...), TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds, CoverageEnabled: pt.CoverageEnabled, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), @@ -514,7 +717,21 @@ 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) { + if len(opts.CompanyIDs) > 0 { + m.companyMu.Lock() + defer m.companyMu.Unlock() + } + pt, err := m.pg.CreateTaskWithOptions(description, goal, opts) if err != nil { return nil, err } @@ -526,6 +743,131 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time return t, nil } +// RenameTaskCategory persists a category name and refreshes every live task DTO +// that references it. taskStateMu keeps this ordered with task reassignment. +func (m *Manager) RenameTaskCategory(id int64, name string) (*pgdb.TaskCategory, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + category, err := m.pg.RenameTaskCategory(id, name) + if err != nil { + return nil, err + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, task := range m.tasks { + task.updateLifecycle(func(state *taskLifecycleState) { + if state.CategoryID != nil && *state.CategoryID == id { + state.CategoryName = category.Name + } + }) + } + return category, nil +} + +// DeleteTaskCategory moves all affected live tasks to the uncategorized bucket. +func (m *Manager) DeleteTaskCategory(id int64) (bool, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + deleted, err := m.pg.DeleteTaskCategory(id) + if err != nil || !deleted { + return deleted, err + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, task := range m.tasks { + task.updateLifecycle(func(state *taskLifecycleState) { + if state.CategoryID != nil && *state.CategoryID == id { + state.CategoryID = nil + state.CategoryName = "" + } + }) + } + return true, nil +} + +// SetTaskCategory updates one live task without interrupting its runtime. +func (m *Manager) SetTaskCategory(taskID string, categoryID *int64) (*pgdb.TaskCategory, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + id, err := strconv.ParseInt(taskID, 10, 64) + if err != nil || id <= 0 { + return nil, pgdb.ErrTaskCategoryTaskNotFound + } + category, err := m.pg.SetTaskCategory(id, categoryID) + if err != nil { + return nil, err + } + m.mu.RLock() + task := m.tasks[taskID] + m.mu.RUnlock() + if task != nil { + task.updateLifecycle(func(state *taskLifecycleState) { + state.CategoryID = cloneInt64Ptr(categoryID) + state.CategoryName = "" + if category != nil { + state.CategoryName = category.Name + } + }) + } + return category, nil +} + +// DeleteCompanyWithAssets keeps the database cascade and live task handles in +// one manager-level critical section. This closes the gap where a task could +// commit its company scope immediately before registration and miss the +// post-delete in-memory sweep. +func (m *Manager) DeleteCompanyWithAssets(id int64, deleteAssets bool) (int64, error) { + m.companyMu.Lock() + defer m.companyMu.Unlock() + assetsDeleted, err := m.pg.Companies().DeleteCompanyWithAssets(id, deleteAssets) + if err != nil { + return 0, err + } + m.mu.Lock() + for _, task := range m.tasks { + task.updateLifecycle(func(state *taskLifecycleState) { + companyIDs := make([]int64, 0, len(state.CompanyIDs)) + for _, companyID := range state.CompanyIDs { + if companyID != id { + companyIDs = append(companyIDs, companyID) + } + } + state.CompanyIDs = companyIDs + }) + } + m.mu.Unlock() + return assetsDeleted, 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.LLMChainRevision, 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() @@ -564,11 +906,215 @@ func (m *Manager) LoadExisting() []*Task { // SetTaskPaused persists a task's paused state. func (m *Manager) SetTaskPaused(id string, paused bool) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() n, err := strconv.ParseInt(id, 10, 64) if err != nil { return err } - return m.pg.SetPaused(n, paused) + if err := m.pg.SetPaused(n, paused); err != nil { + return err + } + m.mu.Lock() + if t := m.tasks[id]; t != nil { + t.updateLifecycle(func(state *taskLifecycleState) { + state.Paused = paused + }) + } + m.mu.Unlock() + return nil +} + +// ApplyTaskAdmission atomically commits the lifecycle fields controlled by the +// concurrency scheduler. Keeping status, paused and queue metadata in one UPDATE +// prevents a failed resume from leaving a task half-revived (for example running +// but still user-paused, or dequeued without an Engine start). +// +// preservePosition applies only when the row is already queued. A repeated +// admission keeps its FIFO timestamp; a task that was explicitly paused and is +// now re-queued receives a fresh tail position. +func (m *Manager) ApplyTaskAdmission(id, expectedStatus, status string, queued bool, mode string, preservePosition bool) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return err + } + if queued { + if mode != "bootstrap" && mode != "resume" { + return fmt.Errorf("invalid queue mode %q", mode) + } + } else { + mode = "" + } + + var queuedAt, completedAt, firstRunAt, deadlineAt sql.NullTime + var committedMode string + err = m.pg.QueryRow(`UPDATE tasks + SET status=$2, + completed_at=CASE + WHEN $2 IN ('done','failed','timeout') THEN COALESCE(completed_at, now()) + ELSE NULL + END, + paused=false, + queued=$3, + queued_at=CASE + WHEN NOT $3 THEN NULL + WHEN $5 AND queued THEN COALESCE(queued_at, now()) + ELSE now() + END, + queue_mode=CASE + WHEN NOT $3 THEN '' + WHEN ($5 AND queued AND queue_mode='bootstrap') OR $4='bootstrap' THEN 'bootstrap' + ELSE 'resume' + END, + first_run_at=CASE + WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL + ELSE first_run_at + END, + deadline_at=CASE + WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL + ELSE deadline_at + END + WHERE id=$1 AND deleted_at IS NULL AND status=$6 + RETURNING queued_at, queue_mode, completed_at, first_run_at, deadline_at`, n, status, queued, mode, preservePosition, expectedStatus). + Scan(&queuedAt, &committedMode, &completedAt, &firstRunAt, &deadlineAt) + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("task %s lifecycle changed before admission (expected status %q)", id, expectedStatus) + } + if err != nil { + return err + } + + m.mu.Lock() + if t := m.tasks[id]; t != nil { + t.updateLifecycle(func(state *taskLifecycleState) { + state.Status = status + state.Paused = false + state.Queued = queued + state.QueueMode = committedMode + state.QueuedAt = 0 + if queuedAt.Valid { + state.QueuedAt = queuedAt.Time.UnixNano() + } + state.CompletedAt = 0 + if completedAt.Valid { + state.CompletedAt = completedAt.Time.Unix() + } + state.FirstRunAt = 0 + if firstRunAt.Valid { + state.FirstRunAt = firstRunAt.Time.Unix() + } + state.DeadlineAt = 0 + if deadlineAt.Valid { + state.DeadlineAt = deadlineAt.Time.Unix() + } + }) + } + m.mu.Unlock() + return nil +} + +// ApplyTaskPause atomically removes a task from the admission queue and records +// the user pause. queue_mode is intentionally retained so resuming a never-run +// bootstrap task still performs goal decomposition, but the next enqueue receives +// a new queued_at timestamp and therefore moves to the FIFO tail. +func (m *Manager) ApplyTaskPause(id string) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return err + } + var mode string + err = m.pg.QueryRow(`UPDATE tasks + SET paused=true, queued=false, queued_at=NULL + WHERE id=$1 AND deleted_at IS NULL AND paused=false + AND status NOT IN ('done','failed','timeout') + RETURNING COALESCE(queue_mode,'')`, n).Scan(&mode) + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("task %s is unavailable for pause", id) + } + if err != nil { + return err + } + m.mu.Lock() + if t := m.tasks[id]; t != nil { + t.updateLifecycle(func(state *taskLifecycleState) { + state.Paused = true + state.Queued = false + state.QueuedAt = 0 + state.QueueMode = mode + }) + } + m.mu.Unlock() + return nil +} + +// EnqueueTask persists the concurrency hold and syncs the in-memory handle. +func (m *Manager) EnqueueTask(id, mode string) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return err + } + if mode != "bootstrap" && mode != "resume" { + return fmt.Errorf("invalid queue mode %q", mode) + } + var queuedAt time.Time + var committedMode string + err = m.pg.QueryRow(`UPDATE tasks + SET queued=true, + queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END, + queue_mode=CASE + WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap' + ELSE 'resume' + END + WHERE id=$1 AND deleted_at IS NULL + RETURNING queued_at, queue_mode`, n, mode).Scan(&queuedAt, &committedMode) + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("task %s is unavailable for enqueue", id) + } + if err != nil { + return err + } + m.mu.Lock() + if t := m.tasks[id]; t != nil { + t.updateLifecycle(func(state *taskLifecycleState) { + state.QueuedAt = queuedAt.UnixNano() + state.Queued = true + state.QueueMode = committedMode + }) + } + m.mu.Unlock() + return nil +} + +// DequeueTask removes the concurrency hold. clearMode=false is used when a user +// pauses a queued task so a later resume still knows whether bootstrap is needed. +func (m *Manager) DequeueTask(id string, clearMode bool) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return err + } + if err := m.pg.Dequeue(n, clearMode); err != nil { + return err + } + m.mu.Lock() + if t := m.tasks[id]; t != nil { + t.updateLifecycle(func(state *taskLifecycleState) { + state.Queued = false + state.QueuedAt = 0 + if clearMode { + state.QueueMode = "" + } + }) + } + m.mu.Unlock() + return nil } // TaskStatus returns a task's current in-memory status (empty if unknown). @@ -576,7 +1122,7 @@ func (m *Manager) TaskStatus(id string) string { m.mu.RLock() defer m.mu.RUnlock() if t := m.tasks[id]; t != nil { - return t.Status + return t.lifecycleSnapshot().Status } return "" } @@ -584,6 +1130,8 @@ func (m *Manager) TaskStatus(id string) string { // StampTaskFirstRun stamps first_run_at + deadline_at on the first real run (idempotent // in DB) and mirrors deadline_at on the live handle. Returns the deadline unix (0 = 不限). func (m *Manager) StampTaskFirstRun(id string) (int64, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() n, err := strconv.ParseInt(id, 10, 64) if err != nil { return 0, err @@ -604,10 +1152,12 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) { } m.mu.Lock() if t := m.tasks[id]; t != nil { - if t.FirstRunAt == 0 { - t.FirstRunAt = time.Now().Unix() - } - t.DeadlineAt = dlUnix + t.updateLifecycle(func(state *taskLifecycleState) { + if state.FirstRunAt == 0 { + state.FirstRunAt = time.Now().Unix() + } + state.DeadlineAt = dlUnix + }) } m.mu.Unlock() return dlUnix, nil @@ -617,6 +1167,8 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) { // (resolves the completed↔timeout race — first terminal writer wins). Reflects the // won status on the live handle. won=false means another terminal already stuck. func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() n, err := strconv.ParseInt(id, 10, 64) if err != nil { return false, err @@ -627,10 +1179,12 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) } m.mu.Lock() if t := m.tasks[id]; t != nil { - t.Status = status - if t.CompletedAt == 0 { - t.CompletedAt = time.Now().Unix() - } + t.updateLifecycle(func(state *taskLifecycleState) { + state.Status = status + if state.CompletedAt == 0 { + state.CompletedAt = time.Now().Unix() + } + }) } m.mu.Unlock() return true, nil @@ -639,6 +1193,8 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) // SetTaskStatus persists a task's lifecycle status (e.g. "done") and reflects it // on the in-memory handle so the derived DTO status shows it without a reload. func (m *Manager) SetTaskStatus(id, status string) error { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() n, err := strconv.ParseInt(id, 10, 64) if err != nil { return err @@ -648,32 +1204,142 @@ func (m *Manager) SetTaskStatus(id, status string) error { } m.mu.Lock() if t := m.tasks[id]; t != nil { - t.Status = status - // mirror the DB's completed_at stamp on the live handle so the DTO shows the - // finish time without a reload (terminal → stamp once; else clear). - if pgdb.IsTerminal(status) { - if t.CompletedAt == 0 { - t.CompletedAt = time.Now().Unix() + t.updateLifecycle(func(state *taskLifecycleState) { + state.Status = status + // Mirror the DB's completed_at stamp on the live handle so the DTO shows + // the finish time without a reload (terminal -> stamp once; else clear). + if pgdb.IsTerminal(status) { + if state.CompletedAt == 0 { + state.CompletedAt = time.Now().Unix() + } + } else { + state.CompletedAt = 0 } - } else { - t.CompletedAt = 0 - } + }) } m.mu.Unlock() 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 { + task.updateLifecycle(func(state *taskLifecycleState) { + kept := make([]int64, 0, len(state.SourceTaskIDs)) + for _, sourceID := range state.SourceTaskIDs { + if sourceID != numericID { + kept = append(kept, sourceID) + } + } + state.SourceTaskIDs = kept + }) + } if m.active == id { m.active = "" for _, t := range m.tasks { @@ -682,7 +1348,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) { diff --git a/server/manager_delete_files_test.go b/server/manager_delete_files_test.go new file mode 100644 index 0000000..48de983 --- /dev/null +++ b/server/manager_delete_files_test.go @@ -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) + } +} diff --git a/server/manager_lifecycle_test.go b/server/manager_lifecycle_test.go new file mode 100644 index 0000000..11be769 --- /dev/null +++ b/server/manager_lifecycle_test.go @@ -0,0 +1,109 @@ +package server + +import ( + "fmt" + "sync" + "testing" +) + +func TestTaskLifecycleSnapshotCopiesContextSlices(t *testing.T) { + task := &Task{ + Status: "running", + SourceTaskIDs: []int64{11, 12}, + CompanyIDs: []int64{21, 22}, + } + + snapshot := task.lifecycleSnapshot() + snapshot.SourceTaskIDs[0] = 99 + snapshot.CompanyIDs[0] = 98 + + fresh := task.lifecycleSnapshot() + if got := fresh.SourceTaskIDs[0]; got != 11 { + t.Fatalf("snapshot source IDs alias task storage: got %d", got) + } + if got := fresh.CompanyIDs[0]; got != 21 { + t.Fatalf("snapshot company IDs alias task storage: got %d", got) + } + + sources := []int64{31, 32} + companies := []int64{41, 42} + task.updateLifecycle(func(state *taskLifecycleState) { + state.SourceTaskIDs = sources + state.CompanyIDs = companies + }) + sources[0] = 97 + companies[0] = 96 + + fresh = task.lifecycleSnapshot() + if got := fresh.SourceTaskIDs[0]; got != 31 { + t.Fatalf("lifecycle update retained caller source slice: got %d", got) + } + if got := fresh.CompanyIDs[0]; got != 41 { + t.Fatalf("lifecycle update retained caller company slice: got %d", got) + } +} + +func TestTaskLifecycleSnapshotConcurrentConsistency(t *testing.T) { + task := &Task{} + writeState := func(generation int64, queued bool) { + task.updateLifecycle(func(state *taskLifecycleState) { + state.Status = map[bool]string{true: "queued", false: "paused"}[queued] + state.Queued = queued + state.Paused = !queued + state.QueuedAt = generation + state.QueueMode = map[bool]string{true: "resume", false: "bootstrap"}[queued] + state.CompletedAt = generation + state.FirstRunAt = generation + state.DeadlineAt = generation + state.SourceTaskIDs = []int64{generation, generation + 1} + state.CompanyIDs = []int64{generation, generation + 1} + }) + } + writeState(1, true) + + var wg sync.WaitGroup + for writer := 0; writer < 4; writer++ { + writer := writer + wg.Add(1) + go func() { + defer wg.Done() + for i := 1; i <= 1000; i++ { + generation := int64(writer*1000 + i) + writeState(generation, generation%2 == 0) + } + }() + } + + errCh := make(chan error, 4) + for reader := 0; reader < 4; reader++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 1000; i++ { + state := task.lifecycleSnapshot() + if state.Queued == state.Paused { + errCh <- fmt.Errorf("queued=%v paused=%v", state.Queued, state.Paused) + return + } + wantStatus := map[bool]string{true: "queued", false: "paused"}[state.Queued] + wantMode := map[bool]string{true: "resume", false: "bootstrap"}[state.Queued] + if state.Status != wantStatus || state.QueueMode != wantMode { + errCh <- fmt.Errorf("inconsistent status/mode: %+v", state) + return + } + generation := state.QueuedAt + if state.CompletedAt != generation || state.FirstRunAt != generation || state.DeadlineAt != generation || + len(state.SourceTaskIDs) != 2 || len(state.CompanyIDs) != 2 || + state.SourceTaskIDs[0] != generation || state.CompanyIDs[0] != generation { + errCh <- fmt.Errorf("mixed lifecycle snapshot: %+v", state) + return + } + } + }() + } + wg.Wait() + close(errCh) + for err := range errCh { + t.Fatal(err) + } +} diff --git a/server/orchestration.go b/server/orchestration.go index 43960ce..3d4b104 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -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) } @@ -301,10 +302,13 @@ func (s *Server) toolPauseTask() actool.CoreTool { TaskID string `json:"task_id"` } _ = json.Unmarshal(in, &a) - if _, ok := s.m.Task(a.TaskID); !ok { + t, ok := s.m.Task(a.TaskID) + if !ok { return actool.Errorf("task 不存在: " + a.TaskID), nil } - s.engine.Pause(a.TaskID) + if _, err := s.applyTaskControlWithCause(t, "pause", agent.AbortPausedByOrchestrator); err != nil { + return actool.Errorf(err.Error()), nil + } return actool.Text("task paused: " + a.TaskID), nil }) } @@ -414,12 +418,13 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool { // deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool. func (s *Server) deriveTaskStatus(t *Task) string { + lifecycle := t.lifecycleSnapshot() switch { - case isTerminalStatus(t.Status): - return t.Status - case t.Paused || s.engine.IsPaused(t.ID): + case isTerminalStatus(lifecycle.Status): + return lifecycle.Status + case lifecycle.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" diff --git a/server/scheduler.go b/server/scheduler.go index 1aeeb17..97cd0b2 100644 --- a/server/scheduler.go +++ b/server/scheduler.go @@ -50,6 +50,7 @@ func (sc *Scheduler) Run(ctx context.Context) { case <-ctx.Done(): return case <-t.C: + sc.s.reconcileConcurrency() // 并发上限:有空位就把排队任务补位启动 sc.step() } } @@ -155,19 +156,22 @@ func (sc *Scheduler) fireFindings(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64) events, err := sc.pg.NewFindingsSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark whether or not any on_finding trigger is active: a + // finding fires only for triggers live at the moment it appears. Otherwise a + // trigger enabled later would replay the entire historical backlog at once. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务发现 finding 触发】\n任务: #%d %s（目标：%s）\n发现: [%s/%s] %s", e.TaskID, e.TaskDesc, e.TaskGoal, e.VulnClass, e.Severity, e.Summary) for _, tr := range want { @@ -185,13 +189,12 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } events, err := sc.pg.MetGoals() if err != nil || len(events) == 0 { return } + // Mark goals as consumed whether or not a trigger is active, so enabling an + // on_goal_met trigger later doesn't replay every already-met goal. fired := sc.firedGoalSet() changed := false for _, e := range events { @@ -200,6 +203,9 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) { } fired[e.NodeID] = true changed = true + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s（目标：%s）\n达成目标: %s", e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary) for _, tr := range want { @@ -220,19 +226,20 @@ func (sc *Scheduler) fireTaskTimeouts(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64) events, err := sc.pg.TimedOutTasksSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s（目标：%s）", e.TaskID, e.TaskDesc, e.TaskGoal) for _, tr := range want { @@ -251,19 +258,20 @@ func (sc *Scheduler) fireTaskCreates(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTaskCreate), 10, 64) events, err := sc.pg.NewTasksSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } msgCtx := fmt.Sprintf("\n\n【本次由任务创建触发】\n任务ID: #%d\n任务描述: %s\n任务目标: %s", e.TaskID, e.TaskDesc, e.TaskGoal) for _, tr := range want { @@ -283,19 +291,20 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) { want = append(want, tr) } } - if len(want) == 0 { - return - } last, _ := strconv.ParseInt(sc.mustState(schedKeyLastToolCall), 10, 64) events, err := sc.pg.NewToolCallsSince(last) if err != nil || len(events) == 0 { return } + // Advance the watermark even with no active trigger — see fireFindings. maxID := last for _, e := range events { if e.NodeID > maxID { maxID = e.NodeID } + if len(want) == 0 { + continue + } errTag := "" if e.ToolIsErr { errTag = "[error] " diff --git a/server/task_control_routes_test.go b/server/task_control_routes_test.go new file mode 100644 index 0000000..26e2ad6 --- /dev/null +++ b/server/task_control_routes_test.go @@ -0,0 +1,46 @@ +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestWorkerControlRoutes(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + + td := t.TempDir() + s := New(context.Background(), m, td, td, td) + h := s.Handler() + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatalf("sign jwt: %v", err) + } + + request := func(path string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"action":"pause"}`)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec + } + + // The single-worker route is still registered and reaches its JSON handler. + single := request("/api/tasks/missing/intents/1/control") + if single.Code != http.StatusNotFound || !strings.Contains(single.Body.String(), `"error":"task not found"`) { + t.Fatalf("single control route unavailable: status=%d body=%s", single.Code, single.Body.String()) + } + + // The former batch route must fall through the mux instead of reaching a task handler. + batch := request("/api/tasks/missing/intents/control/batch") + if batch.Code != http.StatusNotFound || !strings.Contains(batch.Body.String(), "404 page not found") { + t.Fatalf("batch control route still registered: status=%d body=%s", batch.Code, batch.Body.String()) + } +} diff --git a/server/task_delete_barrier_test.go b/server/task_delete_barrier_test.go new file mode 100644 index 0000000..361b161 --- /dev/null +++ b/server/task_delete_barrier_test.go @@ -0,0 +1,364 @@ +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, false) + 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", true) + if !e.IsPaused("42") { + t.Fatal("aborted delete must preserve a pre-existing user pause") + } +} + +func TestBeginTaskDeleteSerializesWithLifecycleLock(t *testing.T) { + s := &Server{engine: NewEngine(nil)} + s.concMu.Lock() + done := make(chan bool, 1) + go func() { done <- s.beginTaskDelete("42") }() + + select { + case <-done: + t.Fatal("delete barrier bypassed the lifecycle lock") + case <-time.After(20 * time.Millisecond): + } + if s.engine.IsDeleting("42") { + t.Fatal("delete barrier was installed before concMu was released") + } + s.concMu.Unlock() + select { + case ok := <-done: + if !ok { + t.Fatal("delete barrier was not installed after lifecycle lock release") + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for delete barrier") + } + s.engine.AbortDelete("42", false) +} + +func TestDeletingRunningTaskKeepsConcurrencySlotUntilCommit(t *testing.T) { + task := &Task{ID: "42", Status: "running"} + s := &Server{engine: NewEngine(nil)} + if !s.engine.BeginDelete(task.ID) { + t.Fatal("install delete barrier") + } + if !s.occupiesConcurrencySlot(task) { + t.Fatal("temporary delete pause released the running task's slot before commit") + } + s.engine.AbortDelete(task.ID, false) +} + +func TestTaskLifecycleRechecksDeleteBarrierAfterConcLock(t *testing.T) { + task := &Task{ID: "42", notify: make(chan struct{}, 1)} + m := &Manager{tasks: map[string]*Task{task.ID: task}} + s := &Server{m: m, engine: NewEngine(m)} + + s.concMu.Lock() + started := make(chan struct{}) + pauseDone := make(chan error, 1) + go func() { + close(started) + _, err := s.applyTaskControl(task, "pause") + pauseDone <- err + }() + <-started + if !s.engine.BeginDelete(task.ID) { + t.Fatal("install delete barrier") + } + s.concMu.Unlock() + if err := <-pauseDone; err == nil || !strings.Contains(err.Error(), "正在删除") { + t.Fatalf("pause passed a delete barrier installed while waiting for concMu: %v", err) + } + s.engine.AbortDelete(task.ID, false) + + s.concMu.Lock() + started = make(chan struct{}) + admitDone := make(chan error, 1) + go func() { + close(started) + _, err := s.admitTask(task, "resume") + admitDone <- err + }() + <-started + if !s.engine.BeginDelete(task.ID) { + t.Fatal("install second delete barrier") + } + s.concMu.Unlock() + if err := <-admitDone; err == nil || !strings.Contains(err.Error(), "being deleted") { + t.Fatalf("admission passed a delete barrier installed while waiting for concMu: %v", err) + } + s.engine.AbortDelete(task.ID, false) +} + +func TestAbortTaskDeleteUsesPersistedPauseAndQueueState(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("delete abort state", "restore persisted barrier", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }() + s := &Server{m: m, engine: NewEngine(m)} + + // An engine-only pause is not authoritative. The committed task is running, + // so aborting deletion must clear this stale in-memory barrier. + s.engine.Pause(task.ID, agent.AbortPausedByUser) + if !s.beginTaskDelete(task.ID) { + t.Fatal("install first delete barrier") + } + s.abortTaskDelete(task.ID) + if s.engine.IsPaused(task.ID) { + t.Fatal("stale engine pause survived although the persisted task is running") + } + + // A queued row must retain its execution barrier even when BeginDelete saw the + // engine as running. This specifically guards against restoring a stale snapshot. + if err := m.EnqueueTask(task.ID, "resume"); err != nil { + t.Fatal(err) + } + if !s.beginTaskDelete(task.ID) { + t.Fatal("install second delete barrier") + } + s.abortTaskDelete(task.ID) + if !s.engine.IsPaused(task.ID) { + t.Fatal("persisted queued task lost its execution barrier after abort") + } +} + +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 TestTaskChatStatusUsesAuthoritativeBusyState(t *testing.T) { + m := &Manager{tasks: map[string]*Task{"7": {ID: "7"}}} + s := &Server{ + m: m, + chatBusy: map[string]bool{"7": true}, + chatCancel: map[string]context.CancelCauseFunc{}, + } + req := httptest.NewRequest(http.MethodGet, "/api/tasks/7/chat/status", nil) + req.SetPathValue("id", "7") + rec := httptest.NewRecorder() + + s.taskChatStatus(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String()) + } + var got struct { + Running bool `json:"running"` + } + if err := json.NewDecoder(rec.Body).Decode(&got); err != nil { + t.Fatal(err) + } + if !got.Running { + t.Fatal("busy main-agent turn was reported idle") + } +} + +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) + } +} diff --git a/traffic/traffic.go b/traffic/traffic.go index 5085aba..e827443 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -11,6 +11,7 @@ import ( "database/sql" "encoding/hex" "encoding/json" + "errors" "fmt" "io" "io/fs" @@ -28,9 +29,9 @@ import ( "time" "unicode/utf8" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" - "github.com/Autumn-27/artex/db" mproxy "github.com/lqqyt2423/go-mitmproxy/proxy" _ "modernc.org/sqlite" ) @@ -307,15 +308,6 @@ VALUES(?,?,?,?,?,?,?)`, return } -<<<<<<< Updated upstream - reqTxt := fmt.Sprintf("%s %s %s\n%s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, headerLines(f.Request.Header), reqBody) - respTxt := fmt.Sprintf("HTTP %d\n%s\n%s", f.Response.StatusCode, headerLines(f.Response.Header), respBody) - _ = os.WriteFile(filepath.Join(exDir, "request.http"), []byte(reqTxt), 0o644) - _ = os.WriteFile(filepath.Join(exDir, "response.http"), []byte(respTxt), 0o644) - meta := fmt.Sprintf(`{"id":%q,"host":%q,"method":%q,"url":%q,"template":%q,"status":%d,"content_type":%q}`, - id, host, method, f.Request.URL.String(), tmpl, f.Response.StatusCode, ct) - _ = os.WriteFile(filepath.Join(exDir, "meta.json"), []byte(meta), 0o644) -======= for _, h := range []string{reqB.hash, respB.hash} { if h == "" { continue @@ -325,7 +317,6 @@ VALUES(?,?,?,?,?,?,?)`, return } } ->>>>>>> Stashed changes if t.fts { // The index is fed from memory, so a body that spilled to _blobs is still @@ -480,6 +471,26 @@ func headerLines(h map[string][]string) string { return b.String() } +// requestHeaderLines restores the HTTP Host header, which net/http stores on +// Request.Host rather than in Header. Keeping it in request.http makes the raw +// capture complete and directly replayable. +func requestHeaderLines(req *mproxy.Request) string { + headers := req.Header.Clone() + headers.Del("Host") + host := req.URL.Host + if raw := req.Raw(); raw != nil && strings.TrimSpace(raw.Host) != "" { + host = raw.Host + } + var b strings.Builder + if host = strings.TrimSpace(host); host != "" { + b.WriteString("Host: ") + b.WriteString(host) + b.WriteByte('\n') + } + b.WriteString(headerLines(headers)) + return b.String() +} + func sanitize(s string) string { r := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "?", "_", "*", "_", "\"", "_", "<", "_", ">", "_", "|", "_") out := r.Replace(s) @@ -944,14 +955,6 @@ func (t *Traffic) reapStage(stageDir string) { // so picking "api.example.com" never sweeps "api.example.com.cn". Returns rows // deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass). func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) { -<<<<<<< Updated upstream - t.wmu.Lock() - defer t.wmu.Unlock() - var deleted int64 - removed := make(map[string]bool) - for _, h := range hosts { - res, err := t.db.Exec(`DELETE FROM exchanges WHERE host=?`, h) -======= stage, err := t.StageDeleteHostsExact(hosts) if err != nil { return 0, err @@ -1028,30 +1031,9 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error stage.tx = tx for _, h := range unique { n, err := t.deleteWhere(tx, `host=?`, h) ->>>>>>> Stashed changes if err != nil { - return deleted, err + return fail(err) } -<<<<<<< Updated upstream - n, _ := res.RowsAffected() - deleted += n - if !removed[h] { - os.RemoveAll(filepath.Join(t.dir, sanitize(h))) - removed[h] = true - } - } - if deleted > 0 { - _ = t.gcBlobs() - } - return deleted, nil -} - -// gcBlobs removes content-addressed blobs in _blobs that no remaining exchange -// tree references. References appear as "@blob sha256:" lines inside the -// tree's request.http/response.http files (see bodyOrBlob). Best-effort: walk -// errors only skip cleanup of the affected files. Callers must hold wmu so this -// never races a concurrent record() writing a fresh blob + tree. -======= stage.deleted += n } @@ -1122,7 +1104,6 @@ func (s *HostDeleteStage) Commit() error { // left the buckets behind permanently. Best-effort: walk errors only skip the // affected entries. Callers must hold wmu so this never races a concurrent // record() writing a fresh blob + reference. ->>>>>>> Stashed changes func (t *Traffic) gcBlobs() error { refs := make(map[string]struct{}) rows, err := t.db.Query(`SELECT DISTINCT hash FROM blob_refs`) diff --git a/traffic/traffic_test.go b/traffic/traffic_test.go index 25d95b9..ea818b5 100644 --- a/traffic/traffic_test.go +++ b/traffic/traffic_test.go @@ -2,12 +2,30 @@ package traffic import ( "fmt" + "net/http" + "net/url" "os" "path/filepath" "strings" "testing" + + mproxy "github.com/lqqyt2423/go-mitmproxy/proxy" ) +func TestRequestHeaderLinesIncludesHost(t *testing.T) { + req := &mproxy.Request{ + URL: &url.URL{Host: "target.example:8443"}, + Header: http.Header{"Accept": []string{"application/json"}}, + } + got := requestHeaderLines(req) + if !strings.Contains(got, "Host: target.example:8443\n") { + t.Fatalf("request headers missing Host: %q", got) + } + if !strings.Contains(got, "Accept: application/json\n") { + t.Fatalf("request headers missing regular header: %q", got) + } +} + // TestDeleteHost verifies the delete contract: rows for hosts containing the // substring are removed together with their file trees, non-matching hosts are // untouched, and the count is right. @@ -162,6 +180,152 @@ VALUES(?,?,?,?,?,?,?,?,?,?,?)`, } } +func TestDeleteHostsExactReportsTreeRemovalFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "api.example.com" + 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", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/1-0001"); err != nil { + t.Fatal(err) + } + + notDir := filepath.Join(dir, "not-a-directory") + if err := os.WriteFile(notDir, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + tr.dir = notDir + + n, err := tr.DeleteHostsExact([]string{host}) + if err == nil { + t.Fatal("DeleteHostsExact returned nil after traffic tree removal failed") + } + if n != 0 { + t.Fatalf("deleted=%d, want 0 after atomic rollback", n) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("rolled-back index count=%d err=%v, want 1", count, err) + } +} + +func TestDeleteHostsExactRollsBackWholeIndexBatch(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + for i, host := range []string{"a.example.com", "b.example.com"} { + id := fmt.Sprintf("1-%04d", i+1) + if err := os.MkdirAll(filepath.Join(dir, host, "GET", id), 0o755); 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(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + } + if _, err := tr.DB().Exec(`CREATE TRIGGER fail_second_host BEFORE DELETE ON exchanges +WHEN OLD.host='b.example.com' BEGIN SELECT RAISE(ABORT, 'forced delete failure'); END`); err != nil { + t.Fatal(err) + } + + if n, err := tr.DeleteHostsExact([]string{"a.example.com", "b.example.com"}); err == nil || n != 0 { + t.Fatalf("DeleteHostsExact failure=(%d,%v), want (0,error)", n, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges`).Scan(&count); err != nil || count != 2 { + t.Fatalf("rolled-back index count=%d err=%v, want 2", count, err) + } + for _, host := range []string{"a.example.com", "b.example.com"} { + if _, err := os.Stat(filepath.Join(dir, host)); err != nil { + t.Fatalf("tree %s changed despite index rollback: %v", host, err) + } + } +} + +func TestStageDeleteHostsExactRollbackRestoresIndexAndTree(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "rollback.example.com" + const id = "1-0001" + tree := filepath.Join(dir, host, "GET", id) + if err := os.MkdirAll(tree, 0o755); err != nil { + t.Fatal(err) + } + marker := filepath.Join(tree, "request.http") + if err := os.WriteFile(marker, []byte("original"), 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(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + if stage.Deleted() != 1 { + t.Fatalf("staged deleted=%d, want 1", stage.Deleted()) + } + if _, err := os.Stat(filepath.Join(dir, host)); !os.IsNotExist(err) { + t.Fatalf("host tree was not staged: %v", err) + } + if err := stage.Rollback(); err != nil { + t.Fatal(err) + } + if got, err := os.ReadFile(marker); err != nil || string(got) != "original" { + t.Fatalf("restored tree content=%q err=%v", got, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("restored index count=%d err=%v, want 1", count, err) + } +} + +func TestStageDeleteHostsExactRollbackReportsRestoreFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "conflict.example.com" + hostDir := filepath.Join(dir, host) + if err := os.MkdirAll(hostDir, 0o755); err != nil { + t.Fatal(err) + } + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + // Simulate an out-of-band conflicting destination. Rollback must surface the + // failed rename instead of claiming the external data was restored. + if err := os.WriteFile(hostDir, []byte("conflict"), 0o644); err != nil { + t.Fatal(err) + } + if err := stage.Rollback(); err == nil || !strings.Contains(err.Error(), "restore") { + t.Fatalf("rollback err=%v, want restore failure", err) + } + if _, err := os.Stat(stage.stageDir); err != nil { + t.Fatalf("staging was removed after failed restore: %v", err) + } +} + // TestDeleteHostGCBlobs verifies blob garbage collection: after a host's trees // are removed, blobs referenced by no remaining exchange are deleted, while // blobs still referenced (including shared ones) survive.