feat(tasks): add runtime controls and inherited context
This commit is contained in:
1 parent
18926c129f
commit
4392b4fee1
50 files changed
+8291
-663
No files matched your search
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
+33
-25
@@ -3,11 +3,13 @@ package agent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"github.com/Autumn-27/artex/db"
|
||||||
"github.com/Autumn-27/norma/agentcore"
|
"github.com/Autumn-27/norma/agentcore"
|
||||||
"github.com/Autumn-27/norma/harness"
|
"github.com/Autumn-27/norma/harness"
|
||||||
"github.com/Autumn-27/norma/llm"
|
"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
|
// 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
|
tkind = kind
|
||||||
tbuf.WriteString(text)
|
tbuf.WriteString(text)
|
||||||
}
|
}
|
||||||
|
lastTool := &runTrace{startedAt: time.Now()}
|
||||||
|
var lastUsage *llm.Usage
|
||||||
|
|
||||||
var finalText string
|
var finalText string
|
||||||
var rerr error
|
var rerr error
|
||||||
for ev, err := range s.Prompt(ctx, input) {
|
for ev, err := range s.Prompt(ctx, input) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
flush()
|
flush()
|
||||||
if ctx.Err() != nil { // ctx cancelled = manual stop, not a failure
|
if ctx.Err() != nil { // engine/user cancellation, not a provider failure
|
||||||
rec(db.Activity{Kind: "result", Summary: "(已手动停止本次运行)"})
|
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()
|
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
|
return finalText, reason, err
|
||||||
}
|
}
|
||||||
switch ev.Kind {
|
switch ev.Kind {
|
||||||
@@ -84,6 +89,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
|
|||||||
flush()
|
flush()
|
||||||
toolNames[ev.ToolUse.ID] = ev.ToolUse.Name
|
toolNames[ev.ToolUse.ID] = ev.ToolUse.Name
|
||||||
in := string(ev.ToolUse.Input)
|
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,
|
rec(db.Activity{Kind: "tool_use", Tool: ev.ToolUse.Name, ToolUseID: ev.ToolUse.ID,
|
||||||
Summary: ev.ToolUse.Name + " " + firstLine(in, 200), Detail: in})
|
Summary: ev.ToolUse.Name + " " + firstLine(in, 200), Detail: in})
|
||||||
case harness.KindToolResult:
|
case harness.KindToolResult:
|
||||||
@@ -92,6 +98,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
|
|||||||
}
|
}
|
||||||
flush()
|
flush()
|
||||||
out := blocksText(ev.ToolResult.Content)
|
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,
|
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})
|
IsError: ev.ToolResult.IsError, Summary: firstLine(out, 200), Detail: out})
|
||||||
case harness.KindText:
|
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.
|
// buffered final-answer text must stay for the KindResult de-dup.
|
||||||
if ev.Usage != nil {
|
if ev.Usage != nil {
|
||||||
u := *ev.Usage
|
u := *ev.Usage
|
||||||
|
lastUsage = &u
|
||||||
rec(db.Activity{Kind: "usage",
|
rec(db.Activity{Kind: "usage",
|
||||||
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
|
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
|
||||||
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
||||||
@@ -120,25 +128,13 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
|
|||||||
tkind = ""
|
tkind = ""
|
||||||
}
|
}
|
||||||
flush() // flush any trailing thinking / non-final text
|
flush() // flush any trailing thinking / non-final text
|
||||||
sum := ev.Terminal.Text
|
sum, detail := ev.Terminal.Text, ev.Terminal.Text
|
||||||
if sum == "" {
|
if sum == "" || ev.Terminal.Reason == harness.ReasonAbortedTools || ev.Terminal.Reason == harness.ReasonAbortedStreaming {
|
||||||
switch ev.Terminal.Reason {
|
sum, detail = terminalText(ctx, ev.Terminal, lastTool)
|
||||||
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) + ")"
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
u := ev.Terminal.Usage // cumulative token usage for this session
|
u := ev.Terminal.Usage // cumulative token usage for this session
|
||||||
rec(db.Activity{Kind: "result", IsError: ev.Terminal.Err != nil,
|
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,
|
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
|
||||||
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
||||||
if ev.Terminal.Err != nil {
|
if ev.Terminal.Err != nil {
|
||||||
@@ -151,6 +147,18 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
|
|||||||
return finalText, reason, rerr
|
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.
|
// blocksText concatenates the text of a tool-result's content blocks.
|
||||||
func blocksText(blocks []llm.ContentBlock) string {
|
func blocksText(blocks []llm.ContentBlock) string {
|
||||||
var b strings.Builder
|
var b strings.Builder
|
||||||
@@ -165,14 +173,14 @@ func blocksText(blocks []llm.ContentBlock) string {
|
|||||||
return b.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 {
|
func firstLine(s string, max int) string {
|
||||||
s = strings.TrimSpace(s)
|
s = strings.TrimSpace(s)
|
||||||
if i := strings.IndexByte(s, '\n'); i >= 0 {
|
if before, _, found := strings.Cut(s, "\n"); found {
|
||||||
s = s[:i]
|
s = before
|
||||||
}
|
}
|
||||||
if len(s) > max {
|
if utf8.RuneCountInString(s) > max {
|
||||||
s = s[:max] + "…"
|
s = string([]rune(s)[:max]) + "…"
|
||||||
}
|
}
|
||||||
return s
|
return s
|
||||||
}
|
}
|
||||||
@@ -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
|
// visible skills/MCP and lets the DB tools table filter/override. DefaultTools
|
||||||
// have no tools-table rows, so they always pass through.
|
// have no tools-table rows, so they always pass through.
|
||||||
base := actool.DefaultTools()
|
base := actool.DefaultTools()
|
||||||
|
ctx = WithRunInfo(ctx, RunInfo{SessionID: sessionID})
|
||||||
tools, def, cleanup := AugmentTools(ctx, agentKey, base)
|
tools, def, cleanup := AugmentTools(ctx, agentKey, base)
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
|
|
||||||
|
|||||||
+18
-3
@@ -7,6 +7,7 @@ import (
|
|||||||
|
|
||||||
"github.com/Autumn-27/artex/db"
|
"github.com/Autumn-27/artex/db"
|
||||||
"github.com/Autumn-27/norma/agentcore"
|
"github.com/Autumn-27/norma/agentcore"
|
||||||
|
"github.com/Autumn-27/norma/llm"
|
||||||
acperm "github.com/Autumn-27/norma/permission"
|
acperm "github.com/Autumn-27/norma/permission"
|
||||||
actool "github.com/Autumn-27/norma/tool"
|
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
|
// no provider is configured or the call yields nothing — the caller then falls
|
||||||
// back to a rule-based split so goal nodes always exist.
|
// 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 数量/交战说明等).
|
// desc is the task's free-text description (背景:靶标范围/flag 数量/交战说明等).
|
||||||
// It is fed alongside the goal so the decomposer no longer splits blind — the
|
// 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.
|
// 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
|
// 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
|
// 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.
|
// 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 {
|
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 {
|
||||||
prov, err := c.NewProvider()
|
if prov == nil {
|
||||||
if err != 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
|
return nil
|
||||||
}
|
}
|
||||||
// worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link
|
// worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link
|
||||||
|
|||||||
+13
-11
@@ -21,6 +21,7 @@ type MainAgent struct {
|
|||||||
model string
|
model string
|
||||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||||
window int // context window in tokens (for compaction)
|
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)
|
maxTurns int // max agent turns per run (0 = unlimited)
|
||||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
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}
|
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
|
// 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).
|
// cert it trusts to verify HTTPS through it (empty addr = direct).
|
||||||
func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert }
|
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.SetAssetStore(as, as.Companies())
|
||||||
}
|
}
|
||||||
tsx.SetTaskID(taskID)
|
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.SetCoverageEnabled(as == nil || as.CoverageEnabled(taskID))
|
||||||
tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced)
|
tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced)
|
||||||
tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running
|
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)。
|
// 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets(不入 prompt)。
|
||||||
base := append(tsx.DropCoverageTools(tsx.MainAgentTools()), actool.DefaultTools()...)
|
base := append(tsx.DropCoverageTools(tsx.MainAgentTools()), actool.DefaultTools()...)
|
||||||
ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
|
ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
|
||||||
>>>>>>> Stashed changes
|
|
||||||
tools, def, cleanup := AugmentTools(ctx, "mainagent", base)
|
tools, def, cleanup := AugmentTools(ctx, "mainagent", base)
|
||||||
tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参
|
tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参
|
||||||
defer cleanup()
|
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
|
BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA
|
||||||
WorkingDir: mainDir, // 本任务工作目录 <workDir>/tasks/<taskID>
|
WorkingDir: mainDir, // 本任务工作目录 <workDir>/tasks/<taskID>
|
||||||
ToolOutputDir: cmdOutDir(mainDir),
|
ToolOutputDir: cmdOutDir(mainDir),
|
||||||
MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management)
|
MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management)
|
||||||
Compaction: compactionConfig(m.window), // long chats stay within the window
|
Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window
|
||||||
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
||||||
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。
|
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。
|
||||||
Settlement: wrapupSettlement("mainagent", nil),
|
Settlement: wrapupSettlement("mainagent", nil),
|
||||||
}
|
}
|
||||||
|
|||||||
+12
-13
@@ -24,6 +24,7 @@ type Planner struct {
|
|||||||
model string
|
model string
|
||||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||||
window int // context window in tokens (for compaction)
|
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)
|
maxTurns int // max agent turns per run (0 = unlimited)
|
||||||
killWork func(intentID int64) error // engine callback to terminate a running work (nil = off)
|
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)
|
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{}}
|
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
|
// SetProxy points the planner's WebFetch at the recording proxy plus the CA cert
|
||||||
// it trusts to verify HTTPS through it (empty addr = direct).
|
// it trusts to verify HTTPS through it (empty addr = direct).
|
||||||
func (p *Planner) SetProxy(addr, caCert string) { p.proxyAddr, p.proxyCACert = addr, caCert }
|
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
|
// 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
|
// the planner looks first at the actual change instead of re-scanning the whole
|
||||||
// overview. Kind:
|
// 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).
|
// "done" — a worker finished intent IntentID (its output conclusion is fetched).
|
||||||
// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要).
|
// "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 文本).
|
// "goal_edited" — the human edited a goal from 总览的目标管理 (OldGoal→NewGoal 文本).
|
||||||
// "cancelled" — the human deleted intent IntentID (Detail = 删除原因). The intent is
|
// "cancelled" — the human deleted intent IntentID (Detail = 删除原因). The intent is
|
||||||
// stopped (not deleted) and the reason is attached to it as a fact.
|
// stopped (not deleted) and the reason is attached to it as a fact.
|
||||||
>>>>>>> Stashed changes
|
|
||||||
type TriggerEvent struct {
|
type TriggerEvent struct {
|
||||||
Kind string
|
Kind string
|
||||||
IntentID int64
|
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)
|
tsx.SetOwnerNode(origin) // planner-side anchors default to the task root (origin fact)
|
||||||
}
|
}
|
||||||
// 领域工具 + 基础默认工具集(Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash)
|
// 领域工具 + 基础默认工具集(Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash)
|
||||||
<<<<<<< Updated upstream
|
|
||||||
base := append(tsx.PlannerTools(), actool.DefaultTools()...)
|
|
||||||
=======
|
|
||||||
// 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets(不入 prompt)。
|
// 资产覆盖度功能关闭时剔除 add_task_scope/list_untested_assets(不入 prompt)。
|
||||||
base := append(tsx.DropCoverageTools(tsx.PlannerTools()), actool.DefaultTools()...)
|
base := append(tsx.DropCoverageTools(tsx.PlannerTools()), actool.DefaultTools()...)
|
||||||
ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
|
ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
|
||||||
>>>>>>> Stashed changes
|
|
||||||
tools, def, cleanup := AugmentTools(ctx, "planner", base)
|
tools, def, cleanup := AugmentTools(ctx, "planner", base)
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
// 关键态势(刚完成的意图 + 预取的完整图)改放【本轮 user 输入】(见下方 input),system
|
// 关键态势(刚完成的意图 + 预取的完整图)改放【本轮 user 输入】(见下方 input),system
|
||||||
@@ -348,7 +347,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
|
|||||||
ToolOutputDir: cmdOutDir(taskDir),
|
ToolOutputDir: cmdOutDir(taskDir),
|
||||||
MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management)
|
MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management)
|
||||||
MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余
|
MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余
|
||||||
Compaction: compactionConfig(p.window),
|
Compaction: compactionConfig(p.compactionWindow()),
|
||||||
// 跨唤醒共享的规划待办:让串行链在多轮之间保留(session 是新的,store 不是)。
|
// 跨唤醒共享的规划待办:让串行链在多轮之间保留(session 是新的,store 不是)。
|
||||||
Todos: p.todoFor(ts.ID()),
|
Todos: p.todoFor(ts.ID()),
|
||||||
// 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、
|
// 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、
|
||||||
@@ -374,7 +373,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
|
|||||||
runCtx := ctx
|
runCtx := ctx
|
||||||
if maxDur > 0 {
|
if maxDur > 0 {
|
||||||
var cancel context.CancelFunc
|
var cancel context.CancelFunc
|
||||||
runCtx, cancel = context.WithTimeout(ctx, maxDur+settleHardGrace)
|
runCtx, cancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
}
|
}
|
||||||
_, _, err = captureRun(runCtx, opts, input,
|
_, _, err = captureRun(runCtx, opts, input,
|
||||||
|
|||||||
@@ -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{}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+329
-148
@@ -23,6 +23,10 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[
|
|||||||
var p map[string]any
|
var p map[string]any
|
||||||
_ = json.Unmarshal(n.Payload, &p)
|
_ = json.Unmarshal(n.Payload, &p)
|
||||||
m := map[string]any{"id": n.ID, "summary": p["summary"], "state": n.State}
|
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
|
// 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
|
// dedup; fall back to legacy payload keys (target_ids plural, then target_id
|
||||||
// single) so intents stored before the rename still surface their anchors.
|
// 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
|
// the model needn't spend a turn calling the tool — every plan round starts with
|
||||||
// an empty context and always needs this first).
|
// an empty context and always needs this first).
|
||||||
func (t *ToolSet) graphOverviewData() map[string]any {
|
func (t *ToolSet) graphOverviewData() map[string]any {
|
||||||
{
|
out := map[string]any{}
|
||||||
out := map[string]any{}
|
// goals summary folded in so the planner needn't call list_goals each round.
|
||||||
// goals summary folded in so the planner needn't call list_goals each round.
|
goals, _ := t.ts.ListByKind(db.KindGoal, 100)
|
||||||
goals, _ := t.ts.ListByKind(db.KindGoal, 100)
|
gsum := make([]map[string]any, 0, len(goals))
|
||||||
gsum := make([]map[string]any, 0, len(goals))
|
for _, g := range goals {
|
||||||
for _, g := range goals {
|
var p map[string]any
|
||||||
var p map[string]any
|
_ = json.Unmarshal(g.Payload, &p)
|
||||||
_ = json.Unmarshal(g.Payload, &p)
|
gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]})
|
||||||
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
|
fr, _ := t.ts.Frontier(100)
|
||||||
// planner reads them every round when generating intents (否则只写不读).
|
out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf)
|
||||||
hints, _ := t.ts.ListByKind(db.KindHint, 50)
|
all, _ := t.ts.ListByKind(db.KindIntent, 300)
|
||||||
hsum := make([]map[string]any, 0, len(hints))
|
var running, recentDone []*db.Node
|
||||||
for _, h := range hints {
|
for _, n := range all {
|
||||||
var p map[string]any
|
switch n.State {
|
||||||
_ = json.Unmarshal(h.Payload, &p)
|
case "running":
|
||||||
hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]})
|
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["running_intents"] = compactIntents(running, parentsOf, yieldsOf)
|
||||||
out["recent_done_intents"] = compactIntents(recentDone, 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
|
e["value"] = r.Domain
|
||||||
case r.Net != "":
|
case r.Net != "":
|
||||||
e["value"] = r.Net
|
e["value"] = r.Net
|
||||||
|
case r.Value != "":
|
||||||
|
e["value"] = r.Value
|
||||||
case r.CompanyID != nil:
|
case r.CompanyID != nil:
|
||||||
e["company_id"] = *r.CompanyID
|
e["company_id"] = *r.CompanyID
|
||||||
if t.cs != nil {
|
if t.cs != nil {
|
||||||
@@ -654,107 +682,163 @@ func (t *ToolSet) relatedTaskOverviews() []map[string]any {
|
|||||||
stats, statsErr := ts.Stats()
|
stats, statsErr := ts.Stats()
|
||||||
|
|
||||||
edges, _ := ts.Edges(5000)
|
edges, _ := ts.Edges(5000)
|
||||||
>>>>>>> Stashed changes
|
|
||||||
parentsOf := map[int64][]int64{}
|
parentsOf := map[int64][]int64{}
|
||||||
yieldsOf := map[int64][]int64{}
|
yieldsOf := map[int64][]int64{}
|
||||||
factFrom := map[int64]int64{}
|
factFrom := map[int64]int64{}
|
||||||
for _, e := range edges {
|
for _, edge := range edges {
|
||||||
switch e.Rel {
|
switch edge.Rel {
|
||||||
case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent)
|
case db.RelDerivedFrom, db.RelSpawns:
|
||||||
parentsOf[e.To] = append(parentsOf[e.To], e.From)
|
parentsOf[edge.To] = append(parentsOf[edge.To], edge.From)
|
||||||
case db.RelYields: // intent --yields--> fact/finding
|
case db.RelYields:
|
||||||
yieldsOf[e.From] = append(yieldsOf[e.From], e.To)
|
yieldsOf[edge.From] = append(yieldsOf[edge.From], edge.To)
|
||||||
factFrom[e.To] = e.From
|
factFrom[edge.To] = edge.From
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
fr, _ := t.ts.Frontier(100)
|
|
||||||
out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf)
|
goals, _ := ts.ListByKind(db.KindGoal, relatedOverviewMaxGoalsPerSource)
|
||||||
all, _ := t.ts.ListByKind(db.KindIntent, 300)
|
goalSummary := make([]map[string]any, 0, len(goals))
|
||||||
var running, recentDone []*db.Node
|
for _, goal := range goals {
|
||||||
for _, n := range all {
|
var payload map[string]any
|
||||||
switch n.State {
|
_ = json.Unmarshal(goal.Payload, &payload)
|
||||||
case "running":
|
goalSummary = append(goalSummary, inheritedMap(map[string]any{
|
||||||
running = append(running, n)
|
"id": goal.ID, "state": goal.State, "text": budget.take(payload["text"], 400),
|
||||||
case "done", "blocked", "exhausted":
|
}, source.Task.TaskID))
|
||||||
if len(recentDone) < 15 {
|
}
|
||||||
recentDone = append(recentDone, n)
|
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)
|
recentFindings := make([]map[string]any, 0, len(findings))
|
||||||
out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf)
|
for _, finding := range findings {
|
||||||
out["frontier_open"] = len(fr)
|
entry := inheritedMap(compactFinding(finding), source.Task.TaskID)
|
||||||
// findings (confirmed vulns) and facts (worker exploration results) are
|
entry["summary"] = budget.take(entry["summary"], 400)
|
||||||
// now distinct node kinds. recent_facts surfaces fact summaries (esp.
|
recentFindings = append(recentFindings, entry)
|
||||||
// negative results) so the planner sees them in one call; full content
|
}
|
||||||
// via node_detail(id).
|
item["recent_findings"] = recentFindings
|
||||||
vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000)
|
recentFacts := make([]map[string]any, 0, len(facts))
|
||||||
factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first
|
for _, fact := range facts {
|
||||||
out["findings"] = len(vulnNodes) // 确认漏洞数(目标判定看它)
|
m := inheritedMap(compactNode(fact), source.Task.TaskID)
|
||||||
out["facts"] = len(factNodes) // 探索事实/结论数(含否定结论)
|
m["summary"] = budget.take(m["summary"], 400)
|
||||||
recentFacts := make([]map[string]any, 0, 20)
|
if from := factFrom[fact.ID]; from > 0 && terminalIntent[from] {
|
||||||
for _, n := range factNodes {
|
m["from_intent"] = from
|
||||||
if len(recentFacts) >= 20 {
|
|
||||||
break
|
|
||||||
}
|
}
|
||||||
m := compactNode(n)
|
var payload map[string]any
|
||||||
if from := factFrom[n.ID]; from > 0 {
|
if json.Unmarshal(fact.Payload, &payload) == nil {
|
||||||
m["from_intent"] = from // 本事实由哪个意图产生
|
if confidence, ok := payload["confidence"].(string); ok && confidence != "" {
|
||||||
}
|
m["confidence"] = confidence
|
||||||
// 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
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
recentFacts = append(recentFacts, m)
|
recentFacts = append(recentFacts, m)
|
||||||
}
|
}
|
||||||
out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?},详情/证据用 node_detail(id)
|
item["recent_facts"] = recentFacts
|
||||||
// the original task (root) so the planner always has it, not just the
|
|
||||||
// decomposed goals.
|
recentDone := recentTerminalIntents(ts, relatedOverviewMaxIntentsPerTask)
|
||||||
if description, goal, err := t.ts.Root(); err == nil {
|
for _, intent := range recentDone {
|
||||||
out["task"] = map[string]any{"description": description, "goal": goal}
|
intent.Inherited = true
|
||||||
|
intent.SourceTaskID = source.Task.TaskID
|
||||||
}
|
}
|
||||||
// coverage:粗略的资产测试覆盖度参考——范围(task_scope)内的资产里,被 fact 碰过的
|
intentResults := compactIntents(recentDone, parentsOf, yieldsOf)
|
||||||
// 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。
|
for i, intent := range recentDone {
|
||||||
if t.as != nil && t.ts != nil && t.taskID > 0 {
|
intentResults[i]["summary"] = budget.take(intentResults[i]["summary"], 400)
|
||||||
if cov, err := t.as.TaskCoverage(t.taskID, t.ts.ID()); err == nil {
|
acts, _, err := ts.ActivityPageForTerminalIntent(intent.ID, 0, 20)
|
||||||
m := map[string]any{
|
if err != nil {
|
||||||
"denominator": cov.Denominator,
|
continue
|
||||||
"tested": cov.Tested,
|
}
|
||||||
"by_type": cov.ByType,
|
var resultSummary, textFallback string
|
||||||
"note": "coverage资产测试覆盖度(包括接口等各种相关资产),粗略估计、仅供参考:容器型资产/大量枚举会让它偏低,勿据此认为已测完;scope 是当前任务测试范围的根资产(root_domain/subdomain/ip/cidr/company 原始行),即覆盖度分母的边界、本任务授权/圈定的目标本身;可用 add_task_scope 增补范围、list_untested_assets 看未测的具体资产【通常不调用list_untested_assets,按照任务推进即可】;",
|
for _, activity := range acts {
|
||||||
|
switch activity.Kind {
|
||||||
|
case "result":
|
||||||
|
resultSummary = activity.Summary
|
||||||
|
case "text":
|
||||||
|
textFallback = activity.Summary
|
||||||
}
|
}
|
||||||
if cov.ScopeRows == 0 {
|
}
|
||||||
m["pct"] = nil
|
if resultSummary == "" {
|
||||||
m["status"] = "范围未锚定"
|
resultSummary = textFallback
|
||||||
} else {
|
}
|
||||||
m["pct"] = cov.Pct
|
if resultSummary != "" {
|
||||||
}
|
intentResults[i]["result_summary"] = budget.take(resultSummary, 800)
|
||||||
// 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
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
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 {
|
func compactNode(n *db.Node) map[string]any {
|
||||||
var p map[string]any
|
var p map[string]any
|
||||||
_ = json.Unmarshal(n.Payload, &p)
|
_ = 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.
|
// 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
|
var p map[string]any
|
||||||
_ = json.Unmarshal(n.Payload, &p)
|
_ = json.Unmarshal(n.Payload, &p)
|
||||||
m := 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)
|
||||||
|
}
|
||||||
if vc, ok := p["vulnclass"]; ok && vc != nil && vc != "" {
|
if vc, ok := p["vulnclass"]; ok && vc != nil && vc != "" {
|
||||||
m["vulnclass"] = vc
|
m["vulnclass"] = vc
|
||||||
}
|
}
|
||||||
@@ -781,16 +872,23 @@ func compactFinding(n *db.Node) map[string]any {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (t *ToolSet) listFindings() actool.CoreTool {
|
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{}),
|
obj(map[string]any{}),
|
||||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||||
f, _ := t.ts.ListByKind(db.KindFinding, 500)
|
f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500)
|
||||||
intentOf, _ := t.ts.FindingIntents() // finding id -> 产生它的 intent id
|
intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id
|
||||||
taskID := t.ts.ID() // exploration id = 任务 id(本 store 内所有 finding 同属)
|
taskID := t.taskID
|
||||||
|
if taskID <= 0 {
|
||||||
|
taskID, _ = t.ts.TaskID()
|
||||||
|
}
|
||||||
out := make([]map[string]any, 0, len(f))
|
out := make([]map[string]any, 0, len(f))
|
||||||
for _, n := range f {
|
for _, n := range f {
|
||||||
m := compactFinding(n)
|
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 {
|
if iid, ok := intentOf[n.ID]; ok {
|
||||||
m["intent_id"] = iid
|
m["intent_id"] = iid
|
||||||
}
|
}
|
||||||
@@ -801,10 +899,10 @@ func (t *ToolSet) listFindings() actool.CoreTool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (t *ToolSet) listFacts() 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{}),
|
obj(map[string]any{}),
|
||||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
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))
|
out := make([]map[string]any, 0, len(f))
|
||||||
for _, n := range f {
|
for _, n := range f {
|
||||||
out = append(out, compactNode(n))
|
out = append(out, compactNode(n))
|
||||||
@@ -814,7 +912,7 @@ func (t *ToolSet) listFacts() actool.CoreTool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (t *ToolSet) nodeDetail() 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"),
|
obj(map[string]any{"id": idp("探索图节点 id(非资产 id)")}, "id"),
|
||||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||||
var a struct {
|
var a struct {
|
||||||
@@ -825,7 +923,7 @@ func (t *ToolSet) nodeDetail() actool.CoreTool {
|
|||||||
if id <= 0 {
|
if id <= 0 {
|
||||||
return actool.Errorf("id 必填"), nil
|
return actool.Errorf("id 必填"), nil
|
||||||
}
|
}
|
||||||
n, err := t.ts.GetNode(id)
|
n, err := t.ts.GetNodeWithSources(id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
@@ -857,9 +955,9 @@ func (t *ToolSet) addOneIntent(it intentItem) (int64, error) {
|
|||||||
// 先校验锚点(建节点前,避免坏锚点留下孤儿意图)。
|
// 先校验锚点(建节点前,避免坏锚点留下孤儿意图)。
|
||||||
parents := pidList(it.ParentIDs)
|
parents := pidList(it.ParentIDs)
|
||||||
for _, pidv := range parents {
|
for _, pidv := range parents {
|
||||||
n, err := t.ts.GetNode(pidv)
|
n, err := t.ts.GetNodeWithSources(pidv)
|
||||||
if err != nil || n == nil {
|
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 {
|
if n.Kind != db.KindFact && n.Kind != db.KindFinding {
|
||||||
return 0, fmt.Errorf("parent_id %d 是 %q 节点,不能作为意图锚点:意图只能锚在已确认的【事实(fact)/发现(finding)】上,不能挂在意图/目标/提示上;顶层全新方向请留空 parent_ids", pidv, n.Kind)
|
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 {
|
if goal == 0 || ev == 0 {
|
||||||
return actool.Errorf("goal_id 和 evidence_id 必填"), nil
|
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.Link(ev, db.RelProves, goal)
|
||||||
_ = t.ts.SetNodeState(goal, "met")
|
_ = t.ts.SetNodeState(goal, "met")
|
||||||
// 每标记一个目标 met,就检查本任务是否【所有目标】都已 met;若是,自动判定
|
// 每标记一个目标 met,就检查本任务是否【所有目标】都已 met;若是,自动判定
|
||||||
@@ -1032,12 +1138,19 @@ func (t *ToolSet) addFinding() actool.CoreTool {
|
|||||||
}
|
}
|
||||||
var id int64
|
var id int64
|
||||||
if t.ts != nil {
|
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
|
var err error
|
||||||
id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors)
|
id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), 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.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)
|
_, _ = 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 != "" {
|
if c := strings.TrimSpace(it.Confidence); c != "" {
|
||||||
payload["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).
|
// 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))
|
id, err := t.ts.AddNode(db.KindFact, payload, 5, "confirmed", t.worker, pidList(it.AssetIDs))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
intent := pid(it.IntentID)
|
|
||||||
if intent <= 0 {
|
|
||||||
intent = defaultIntent
|
|
||||||
}
|
|
||||||
if intent > 0 {
|
if intent > 0 {
|
||||||
_ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> fact
|
_ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> fact
|
||||||
}
|
}
|
||||||
@@ -1422,6 +1541,10 @@ func (t *ToolSet) killWorkTool() actool.CoreTool {
|
|||||||
if id <= 0 {
|
if id <= 0 {
|
||||||
return actool.Errorf("intent_id 必填"), nil
|
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 {
|
if err := t.killWork(id); err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
@@ -1452,6 +1575,10 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool {
|
|||||||
if id <= 0 {
|
if id <= 0 {
|
||||||
return actool.Errorf("intent_id 必填"), nil
|
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) == "" {
|
if strings.TrimSpace(a.Message) == "" {
|
||||||
return actool.Errorf("message 必填"), nil
|
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.
|
// getWorkerOutput returns a work's final (or截至中止时的) conclusion text by intent id.
|
||||||
func (t *ToolSet) getWorkerOutput() actool.CoreTool {
|
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"),
|
obj(map[string]any{"intent_id": idp("意图 id(= work 句柄)")}, "intent_id"),
|
||||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||||
var a struct {
|
var a struct {
|
||||||
@@ -1475,7 +1602,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool {
|
|||||||
if id <= 0 {
|
if id <= 0 {
|
||||||
return actool.Errorf("intent_id 必填"), nil
|
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 {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
@@ -1494,16 +1628,25 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool {
|
|||||||
pick, terminated = fallback, true
|
pick, terminated = fallback, true
|
||||||
}
|
}
|
||||||
if pick == nil {
|
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
|
return actool.Text("(该 work 尚无任何输出)"), nil
|
||||||
}
|
}
|
||||||
detail, _ := t.ts.ActivityDetail(pick.ID)
|
detail, _ := t.ts.ActivityDetailWithSources(pick.ID)
|
||||||
if detail == "" {
|
if detail == "" {
|
||||||
detail = pick.Summary
|
detail = pick.Summary
|
||||||
}
|
}
|
||||||
return jsonResult(map[string]any{
|
result := map[string]any{
|
||||||
"intent_id": id, "worker_name": pick.Worker, "final_text": detail,
|
"intent_id": id, "worker_name": pick.Worker, "final_text": detail,
|
||||||
"summary": pick.Summary, "is_error": pick.IsError, "terminated": terminated,
|
"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 {
|
func traceSteps(acts []db.Activity) []map[string]any {
|
||||||
steps := make([]map[string]any, 0, len(acts))
|
steps := make([]map[string]any, 0, len(acts))
|
||||||
for i := range 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,
|
"step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool,
|
||||||
"is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100),
|
"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
|
return steps
|
||||||
}
|
}
|
||||||
@@ -1530,7 +1677,7 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool {
|
|||||||
"① 只传 intent_id → 返回该 work 每一步的摘要流(summary≤100字,含 step_id;只是动作轮廓,不含完整输出);\n"+
|
"① 只传 intent_id → 返回该 work 每一步的摘要流(summary≤100字,含 step_id;只是动作轮廓,不含完整输出);\n"+
|
||||||
"② intent_id + q → 只返回命中关键字的步骤摘要(在摘要和完整输出里都搜;仍只给 summary,要看内容用③);\n"+
|
"② intent_id + q → 只返回命中关键字的步骤摘要(在摘要和完整输出里都搜;仍只给 summary,要看内容用③);\n"+
|
||||||
"③ intent_id + step_ids(最多5个)→ 返回这些步骤的完整内容(detail)。\n"+
|
"③ intent_id + step_ids(最多5个)→ 返回这些步骤的完整内容(detail)。\n"+
|
||||||
"典型流程:先①/②定位可疑步骤的 step_id,再用③取其完整输出。不含思考(thinking)步骤。",
|
"典型流程:先①/②定位可疑步骤的 step_id,再用③取其完整输出。不含思考(thinking)步骤。支持直接关联任务的历史 trace;其结果带 source_task_id/inherited=true 且只读。",
|
||||||
obj(map[string]any{
|
obj(map[string]any{
|
||||||
"intent_id": idp("意图 id(= work 句柄)"),
|
"intent_id": idp("意图 id(= work 句柄)"),
|
||||||
"q": str("关键字:只返回摘要/完整输出命中它的步骤(可选;与 step_ids 互斥)"),
|
"q": str("关键字:只返回摘要/完整输出命中它的步骤(可选;与 step_ids 互斥)"),
|
||||||
@@ -1549,6 +1696,13 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool {
|
|||||||
if id <= 0 {
|
if id <= 0 {
|
||||||
return actool.Errorf("intent_id 必填"), nil
|
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.
|
// ③ detail drill-down by step ids (max 5), thinking excluded by the store.
|
||||||
if len(a.StepIDs) > 0 {
|
if len(a.StepIDs) > 0 {
|
||||||
var ids []int64
|
var ids []int64
|
||||||
@@ -1560,31 +1714,47 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool {
|
|||||||
if len(ids) > 5 {
|
if len(ids) > 5 {
|
||||||
return actool.Errorf("step_ids 一次最多 5 个(先用不带 step_ids 的调用定位,再分批取)"), nil
|
return actool.Errorf("step_ids 一次最多 5 个(先用不带 step_ids 的调用定位,再分批取)"), nil
|
||||||
}
|
}
|
||||||
acts, err := t.ts.ActivityByIDs(ids)
|
acts, err := t.ts.ActivityByIDsWithSources(ids)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
steps := make([]map[string]any, 0, len(acts))
|
steps := make([]map[string]any, 0, len(acts))
|
||||||
for i := range 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,
|
"step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool,
|
||||||
"is_error": acts[i].IsError, "detail": acts[i].Detail,
|
"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.
|
// ①/② summary stream, optionally keyword-filtered; 100-char summaries.
|
||||||
var acts []db.Activity
|
var acts []db.Activity
|
||||||
var err error
|
var err error
|
||||||
if strings.TrimSpace(a.Q) != "" {
|
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 {
|
} else {
|
||||||
acts, err = t.ts.ActivityTrace(id, a.Limit)
|
acts, err = t.ts.ActivityTraceWithSources(id, a.Limit)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), 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
|
return actool.Errorf("q 必填"), nil
|
||||||
}
|
}
|
||||||
// 排除调用者自身这条意图的步骤(worker 的自有 trace 已在其上下文里)。
|
// 排除调用者自身这条意图的步骤(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 {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
@@ -1620,11 +1790,15 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool {
|
|||||||
if acts[i].NodeID != nil {
|
if acts[i].NodeID != nil {
|
||||||
intent = *acts[i].NodeID
|
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,
|
"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,
|
"kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError,
|
||||||
"summary": firstLine(acts[i].Summary, 100),
|
"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})
|
return jsonResult(map[string]any{"query": a.Q, "hits": hits})
|
||||||
})
|
})
|
||||||
@@ -1655,13 +1829,16 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool {
|
|||||||
if limit <= 0 {
|
if limit <= 0 {
|
||||||
limit = 50
|
limit = 50
|
||||||
}
|
}
|
||||||
all, err := t.ts.ListByKind(db.KindIntent, 500)
|
all, err := t.ts.ListByKindWithSources(db.KindIntent, 500)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
q := strings.ToLower(strings.TrimSpace(a.Q))
|
q := strings.ToLower(strings.TrimSpace(a.Q))
|
||||||
out := make([]map[string]any, 0, limit)
|
out := make([]map[string]any, 0, limit)
|
||||||
for _, n := range all {
|
for _, n := range all {
|
||||||
|
if n.Inherited && n.State == "running" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
switch n.State {
|
switch n.State {
|
||||||
case "running", "done", "exhausted", "blocked", "stopped": // has run → has a process
|
case "running", "done", "exhausted", "blocked", "stopped": // has run → has a process
|
||||||
default:
|
default:
|
||||||
@@ -1673,7 +1850,11 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool {
|
|||||||
if q != "" && !strings.Contains(strings.ToLower(summary), q) {
|
if q != "" && !strings.Contains(strings.ToLower(summary), q) {
|
||||||
continue
|
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 {
|
if len(out) >= limit {
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|||||||
+28
-16
@@ -3,6 +3,7 @@ package agent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
@@ -284,6 +285,18 @@ func (t *ToolSet) insertAssets() actool.CoreTool {
|
|||||||
results = append(results, result{Index: i, ID: id, Type: typ})
|
results = append(results, result{Index: i, ID: id, Type: typ})
|
||||||
t.writes.Assets++
|
t.writes.Assets++
|
||||||
t.anchorOwner(id)
|
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 顶层显式插入的这一项,按其
|
// 自动入测试范围(source='auto'):只对 worker 顶层显式插入的这一项,按其
|
||||||
// 类型加保守范围;side-effect 派生的资产不经此处,故范围不盲目扩大。taskID=0 时无操作。
|
// 类型加保守范围;side-effect 派生的资产不经此处,故范围不盲目扩大。taskID=0 时无操作。
|
||||||
// 资产覆盖度功能关闭时不再累积测试范围(分母);related=false(与当前任务无关)
|
// 资产覆盖度功能关闭时不再累积测试范围(分母);related=false(与当前任务无关)
|
||||||
@@ -310,14 +323,14 @@ func (t *ToolSet) insertAssets() actool.CoreTool {
|
|||||||
func (t *ToolSet) addCompanyScope() actool.CoreTool {
|
func (t *ToolSet) addCompanyScope() actool.CoreTool {
|
||||||
return writeTool(
|
return writeTool(
|
||||||
"add_company_scope",
|
"add_company_scope",
|
||||||
"把根域名/IP/IP段(CIDR)加入某公司的【资产范围】——加入后系统自动认领所有命中的资产(现有+未来)。\n"+
|
"把域名/IP/CIDR/ICP备案/企业关键词加入某公司的【资产范围】——域名、网络和ICP会自动认领命中的资产,关键词只提供给Agent作为范围提示。\n"+
|
||||||
"公司名唯一:company 不存在则新建,已存在则复用(只把范围并进去)。\n"+
|
"公司名唯一:company 不存在则新建,已存在则复用(只把范围并进去)。\n"+
|
||||||
"scope 一行一条,支持:根域名 / 单个 IP / CIDR 网段(如 1.2.3.0/24)。\n"+
|
"scope 一行一条,系统自动识别:根域名 / URL / 单个 IP / CIDR 网段 / ICP备案 / 企业关键词。\n"+
|
||||||
"务必给 reason 说明归属依据(whois/证书/ASN 等)。\n"+
|
"务必给 reason 说明归属依据(whois/证书/ASN 等)。\n"+
|
||||||
"护栏:拒绝裸 TLD 与过宽网段(IPv4<=/16、IPv6<=/32),非法行会被跳过并在 errors 返回。",
|
"护栏:拒绝裸 TLD 与过宽网段(IPv4前缀需为/16-/32、IPv6前缀需为/32-/128),非法行会被跳过并在 errors 返回。",
|
||||||
obj(map[string]any{
|
obj(map[string]any{
|
||||||
"company": str("公司名(不存在则新建、存在则复用;名称唯一)"),
|
"company": str("公司名(不存在则新建、存在则复用;名称唯一)"),
|
||||||
"scope": str("资产范围,一行一条:根域名 / IP / CIDR"),
|
"scope": str("资产范围,一行一条:域名 / URL / IP / CIDR / ICP备案 / 企业关键词"),
|
||||||
"reason": str("归属依据(证据/来源),务必填写"),
|
"reason": str("归属依据(证据/来源),务必填写"),
|
||||||
"logo": str("公司图标 URL(可选;仅新建公司时生效)"),
|
"logo": str("公司图标 URL(可选;仅新建公司时生效)"),
|
||||||
}, "company", "scope"),
|
}, "company", "scope"),
|
||||||
@@ -365,14 +378,14 @@ func (t *ToolSet) addTaskScope() actool.CoreTool {
|
|||||||
return writeTool(
|
return writeTool(
|
||||||
"add_task_scope",
|
"add_task_scope",
|
||||||
"把测试范围加入【本任务】——这是资产测试覆盖度的分母,也是本任务的授权边界。\n"+
|
"把测试范围加入【本任务】——这是资产测试覆盖度的分母,也是本任务的授权边界。\n"+
|
||||||
"kind 支持:company(整个公司名下资产) / root_domain(整个根域,含所有子域) / subdomain(单个精确子域) / ip / cidr。\n"+
|
"kind 支持:company(整个公司名下资产) / root_domain(整个根域,含所有子域) / subdomain(单个精确子域) / ip / cidr / icp / keyword。\n"+
|
||||||
"说明:worker 逐个碰到的主机会被系统【自动】加进范围(精确子域);本工具用于【主动扩大】——把整个根域/整个公司纳入,或补充指定某子域/IP。\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 数组。",
|
"务必给 reason 说明依据(可审计)。多条用 entries 数组。",
|
||||||
obj(map[string]any{
|
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"}},
|
"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"),
|
"kind": str("[单条] company / root_domain / subdomain / ip / cidr / icp / keyword"),
|
||||||
"value": str("[单条] 公司名或id / 域名 / IP / CIDR"),
|
"value": str("[单条] 公司名或id / 域名 / IP / CIDR / ICP / 关键词"),
|
||||||
"reason": str("加入依据(用于审计),务必填写"),
|
"reason": str("加入依据(用于审计),务必填写"),
|
||||||
}),
|
}),
|
||||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||||
@@ -399,12 +412,12 @@ func (t *ToolSet) addTaskScope() actool.CoreTool {
|
|||||||
var added []map[string]any
|
var added []map[string]any
|
||||||
errs := map[string]string{}
|
errs := map[string]string{}
|
||||||
for i, e := range items {
|
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 {
|
if err != nil {
|
||||||
errs[strconv.Itoa(i)] = err.Error()
|
errs[strconv.Itoa(i)] = err.Error()
|
||||||
continue
|
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}
|
out := map[string]any{"added": added}
|
||||||
if len(errs) > 0 {
|
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-
|
// listUntestedAssets lets the plan agent pull the current + directly inherited
|
||||||
// tested assets on demand (filter by type, paginated) — so it can decide what to
|
// scope's not-yet-tested assets on demand (filter by type, paginated).
|
||||||
// test next itself, instead of coverage pushing a backlog list into every prompt.
|
|
||||||
func (t *ToolSet) listUntestedAssets() actool.CoreTool {
|
func (t *ToolSet) listUntestedAssets() actool.CoreTool {
|
||||||
return readTool(
|
return readTool(
|
||||||
"list_untested_assets",
|
"list_untested_assets",
|
||||||
"查询【本任务】范围内、还没被测过的资产(供你自己判断要不要补测,不代替你决策)。\n"+
|
"查询【本任务及直接关联任务】范围内、还没被事实锚点覆盖的资产(关联范围只读,供你自己判断要不要补测,不代替你决策)。\n"+
|
||||||
"可选按资产类型过滤:root_domain/subdomain/service/app/endpoint/ip。\n"+
|
"可选按资产类型过滤:root_domain/subdomain/service/app/endpoint/ip。\n"+
|
||||||
"分页:page 从 1 起、page_size 默认 10。返回 {assets:[{id,type,label}], total, page, page_size}。仅任务上下文可用。",
|
"分页:page 从 1 起、page_size 默认 10。返回 {assets:[{id,type,label}], total, page, page_size}。仅任务上下文可用。",
|
||||||
obj(map[string]any{
|
obj(map[string]any{
|
||||||
@@ -449,7 +461,7 @@ func (t *ToolSet) listUntestedAssets() actool.CoreTool {
|
|||||||
a.PageSize = 10
|
a.PageSize = 10
|
||||||
}
|
}
|
||||||
offset := (a.Page - 1) * a.PageSize
|
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 {
|
if err != nil {
|
||||||
return actool.Errorf(err.Error()), nil
|
return actool.Errorf(err.Error()), nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+26
-3
@@ -49,6 +49,7 @@ type Worker struct {
|
|||||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||||
window int // context window in tokens (for compaction)
|
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)
|
maxTurns int // max agent turns per run (0 = unlimited)
|
||||||
// runTimeout is the wall-clock budget for the main exploration of one intent
|
// 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
|
// (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}
|
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
|
// SetProxy configures the recording proxy address that workers route target
|
||||||
// traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through
|
// traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through
|
||||||
// that MITM proxy. Empty addr disables the hint.
|
// 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).
|
// phase, Bash is hidden via Settlement.DisabledTools (no local gating needed).
|
||||||
base := append(tsx.WorkerTools(), w.extraTools...)
|
base := append(tsx.WorkerTools(), w.extraTools...)
|
||||||
base = append(base, actool.DefaultTools()...)
|
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, def, cleanup := AugmentTools(ctx, "worker", base)
|
||||||
tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参
|
tools = tsx.StripCoverageParams(tools) // 覆盖度关闭时隐藏 insert_assets 的 related 入参
|
||||||
defer cleanup()
|
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);
|
// large tool output spills to cmd-output/ with a head + pointer (SDK tool.Capture);
|
||||||
// full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。
|
// full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。
|
||||||
ToolOutputDir: cmdOutDir(runDir),
|
ToolOutputDir: cmdOutDir(runDir),
|
||||||
Compaction: compactionConfig(w.window), // long tool-heavy runs stay within the window
|
Compaction: compactionConfig(w.compactionWindow()), // long tool-heavy runs stay within the window
|
||||||
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
||||||
}
|
}
|
||||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||||
opts.Hooks = hooks
|
opts.Hooks = hooks
|
||||||
@@ -371,6 +382,18 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db.
|
|||||||
s := agentcore.NewSession(opts)
|
s := agentcore.NewSession(opts)
|
||||||
defer s.Close() // release the session's background-task manager (temp dir + processes)
|
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):
|
// Budgets + settlement are owned by the SDK (MaxTurns/MaxDuration + Settlement):
|
||||||
// on hit it runs a wrap-up turn and finishes with ReasonMaxTurns/ReasonTimeout.
|
// 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
|
// 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
|
runCtx := ctx
|
||||||
if maxDur > 0 {
|
if maxDur > 0 {
|
||||||
var runCancel context.CancelFunc
|
var runCancel context.CancelFunc
|
||||||
runCtx, runCancel = context.WithTimeout(ctx, maxDur+settleHardGrace)
|
runCtx, runCancel = context.WithTimeoutCause(ctx, maxDur+settleHardGrace, AbortRunHardTimeout)
|
||||||
defer runCancel()
|
defer runCancel()
|
||||||
}
|
}
|
||||||
_, reason, err := captureRunSession(runCtx, s, input, emitWrap)
|
_, reason, err := captureRunSession(runCtx, s, input, emitWrap)
|
||||||
|
|||||||
+19
-1
@@ -16,6 +16,7 @@ import (
|
|||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/Autumn-27/artex/agent"
|
||||||
"github.com/Autumn-27/artex/config"
|
"github.com/Autumn-27/artex/config"
|
||||||
"github.com/Autumn-27/artex/server"
|
"github.com/Autumn-27/artex/server"
|
||||||
)
|
)
|
||||||
@@ -70,8 +71,10 @@ func main() {
|
|||||||
log.Printf("[config] 配置文件: %s (不存在 — 将仅尝试环境变量 ARTEX_PG_DSN)", cfgPath)
|
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()
|
defer stop()
|
||||||
|
ctx, shutdown := shutdownContext(sigCtx)
|
||||||
|
defer shutdown(agent.AbortShutdown)
|
||||||
|
|
||||||
mgr, err := server.NewManager(*dataDir, *proxy)
|
mgr, err := server.NewManager(*dataDir, *proxy)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -104,3 +107,18 @@ func main() {
|
|||||||
defer cancel()
|
defer cancel()
|
||||||
_ = httpSrv.Shutdown(shutdownCtx)
|
_ = 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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -103,7 +103,7 @@ type builtinAgent struct {
|
|||||||
key, name, role, desc string
|
key, name, role, desc string
|
||||||
vars []promptVar
|
vars []promptVar
|
||||||
interactiveShell bool // 建行时的默认交互式 shell 开关;ON CONFLICT 不覆盖用户后续手动开关
|
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 }
|
type promptVar struct{ name, desc, example, source string }
|
||||||
@@ -145,7 +145,7 @@ func (d *DB) seedBuiltins() error {
|
|||||||
var agentID int64
|
var agentID int64
|
||||||
err := d.QueryRow(`
|
err := d.QueryRow(`
|
||||||
INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds)
|
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
|
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)
|
RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
+314
-52
@@ -3,11 +3,17 @@ package db
|
|||||||
import (
|
import (
|
||||||
"database/sql"
|
"database/sql"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"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
|
// 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.
|
// 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
|
// 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.
|
// Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint.
|
||||||
type Node struct {
|
type Node struct {
|
||||||
ID int64 `json:"id"`
|
ID int64 `json:"id"`
|
||||||
Kind string `json:"kind"`
|
Kind string `json:"kind"`
|
||||||
Payload json.RawMessage `json:"payload"`
|
Payload json.RawMessage `json:"payload"`
|
||||||
Priority int `json:"priority"`
|
Priority int `json:"priority"`
|
||||||
State string `json:"state"`
|
State string `json:"state"`
|
||||||
Origin string `json:"origin,omitempty"`
|
Origin string `json:"origin,omitempty"`
|
||||||
Owner string `json:"owner,omitempty"`
|
Owner string `json:"owner,omitempty"`
|
||||||
Anchors []int64 `json:"anchors,omitempty"`
|
BlockedReason string `json:"blocked_reason,omitempty"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
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).
|
// Activity is one worker execution step (= old task_activity).
|
||||||
type Activity struct {
|
type Activity struct {
|
||||||
ID int64 `json:"id"`
|
ID int64 `json:"id"`
|
||||||
NodeID *int64 `json:"node_id,omitempty"`
|
NodeID *int64 `json:"node_id,omitempty"`
|
||||||
Worker string `json:"worker,omitempty"`
|
Worker string `json:"worker,omitempty"`
|
||||||
Kind string `json:"kind,omitempty"`
|
Kind string `json:"kind,omitempty"`
|
||||||
Tool string `json:"tool,omitempty"`
|
Tool string `json:"tool,omitempty"`
|
||||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||||
IsError bool `json:"is_error"`
|
IsError bool `json:"is_error"`
|
||||||
Summary string `json:"summary,omitempty"`
|
Summary string `json:"summary,omitempty"`
|
||||||
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
|
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
|
||||||
CreatedAt time.Time `json:"created_at"`
|
Metadata json.RawMessage `json:"metadata,omitempty"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
// Token usage — set only on the terminal kind='result' record; nil elsewhere.
|
// Token usage — set only on the terminal kind='result' record; nil elsewhere.
|
||||||
InputTokens *int `json:"input_tokens,omitempty"`
|
InputTokens *int `json:"input_tokens,omitempty"`
|
||||||
OutputTokens *int `json:"output_tokens,omitempty"`
|
OutputTokens *int `json:"output_tokens,omitempty"`
|
||||||
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
|
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
|
||||||
CacheWriteTokens *int `json:"cache_write_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).
|
// TokenUsage is a per-worker token aggregate (TokenStatsByWorker).
|
||||||
@@ -64,6 +76,17 @@ type TokenUsage struct {
|
|||||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
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.
|
// DailyTokenBucket is one day's global token aggregate across all tasks.
|
||||||
type DailyTokenBucket struct {
|
type DailyTokenBucket struct {
|
||||||
Day string `json:"day"` // "YYYY-MM-DD"
|
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).
|
// SetNodeState updates any node's state (never deletes).
|
||||||
func (s *ExplorationStore) SetNodeState(id int64, state string) error {
|
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
|
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);
|
// 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
|
// 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.
|
// 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.
|
// SetIntentState updates an intent node's state.
|
||||||
>>>>>>> Stashed changes
|
|
||||||
// ResetRunningIntents returns any intent left in 'running' back to 'open' so it
|
// 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
|
// 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
|
// 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) {
|
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)
|
WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
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)
|
// 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
|
// back to 'open' so a worker re-claims it (graph writes it already produced stay).
|
||||||
// already produced stay). done/open/running are left untouched. Returns whether a row
|
// The worker resumes from its prior LLM transcript instead of restarting from scratch.
|
||||||
// changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button.
|
// done/open/running are left untouched. Returns whether a row changed. Used by "重跑".
|
||||||
func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) {
|
func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) {
|
||||||
res, err := s.db.Exec(`UPDATE exploration_nodes
|
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'
|
WHERE id=$1 AND exploration_id=$2 AND kind='intent'
|
||||||
AND state IN ('blocked','exhausted','stopped')`, id, s.expID)
|
AND state IN ('blocked','exhausted','stopped')`, id, s.expID)
|
||||||
if err != nil {
|
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'
|
// 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
|
// (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) {
|
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)
|
WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
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 states stamp completed_at; reopening (back to open/running) clears it.
|
||||||
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
|
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
|
||||||
_, err := s.db.Exec(`UPDATE exploration_nodes
|
_, 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)
|
WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal)
|
||||||
return err
|
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
|
// CompareAndSetIntentState transitions one local intent only when its current
|
||||||
// state still matches expected. It is the state boundary used by pause/resume and
|
// 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.
|
// 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`
|
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) {
|
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
|
||||||
var n Node
|
var n Node
|
||||||
var payload []byte
|
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
|
return nil, err
|
||||||
}
|
}
|
||||||
n.Payload = json.RawMessage(payload)
|
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()
|
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).
|
// Frontier returns open intents ordered by priority desc then id (FIFO tiebreak).
|
||||||
func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) {
|
func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) {
|
||||||
if limit <= 0 {
|
if limit <= 0 {
|
||||||
@@ -851,12 +891,16 @@ func (s *ExplorationStore) Stats() (map[string]int, error) {
|
|||||||
|
|
||||||
// AppendActivity records one worker step and returns its id.
|
// AppendActivity records one worker step and returns its id.
|
||||||
func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) {
|
func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) {
|
||||||
|
metadata := a.Metadata
|
||||||
|
if len(metadata) == 0 {
|
||||||
|
metadata = json.RawMessage(`{}`)
|
||||||
|
}
|
||||||
var id int64
|
var id int64
|
||||||
err := s.db.QueryRow(`
|
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)
|
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)
|
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),
|
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
|
return id, err
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -986,6 +1030,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
|
|||||||
return out, rows.Err()
|
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.
|
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
|
||||||
// Returns items and the new cursor (max id seen).
|
// Returns items and the new cursor (max id seen).
|
||||||
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
|
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 rows *sql.Rows
|
||||||
var err error
|
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 {
|
if nodeID != nil {
|
||||||
rows, err = s.db.Query(`SELECT `+cols+`
|
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)
|
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()
|
defer rows.Close()
|
||||||
out := []Activity{}
|
out := []Activity{}
|
||||||
cursor := sinceID
|
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() {
|
for rows.Next() {
|
||||||
var a Activity
|
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,
|
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 {
|
if limit <= 0 {
|
||||||
limit = 200
|
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}
|
args := []any{s.expID}
|
||||||
cond, cargs := f.cond(len(args) + 1)
|
cond, cargs := f.cond(len(args) + 1)
|
||||||
args = append(args, cargs...)
|
args = append(args, cargs...)
|
||||||
@@ -1071,6 +1187,51 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg)
|
|||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
desc := []Activity{}
|
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() {
|
for rows.Next() {
|
||||||
var a Activity
|
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,
|
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 {
|
if hasMore {
|
||||||
desc = desc[:limit]
|
desc = desc[:limit]
|
||||||
}
|
}
|
||||||
// reverse the newest-first window into ascending id order for display.
|
|
||||||
out := make([]Activity, len(desc))
|
out := make([]Activity, len(desc))
|
||||||
for i, a := range desc {
|
for i, a := range desc {
|
||||||
out[len(desc)-1-i] = a
|
out[len(desc)-1-i] = a
|
||||||
@@ -1152,6 +1312,26 @@ ORDER BY id LIMIT $3`, s.expID, nodeID, limit)
|
|||||||
return scanTrace(rows)
|
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
|
// ActivityTraceSearch finds steps whose summary OR detail matches q
|
||||||
// (case-insensitive), returning summaries only. nodeID != nil scopes to one work;
|
// (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
|
// 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)
|
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
|
// AssetRef is a compact exploration node (intent/fact/finding) that references an
|
||||||
// asset via an anchor — for the coverage graph's node drawer.
|
// asset via an anchor — for the coverage graph's node drawer.
|
||||||
type AssetRef struct {
|
type AssetRef struct {
|
||||||
ID int64 `json:"id"`
|
ID int64 `json:"id"`
|
||||||
Kind string `json:"kind"`
|
Kind string `json:"kind"`
|
||||||
State string `json:"state"`
|
State string `json:"state"`
|
||||||
Summary string `json:"summary"`
|
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
|
// 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()
|
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
|
// AddStandaloneFinding writes a finding to the standalone findings table, which
|
||||||
// persists across task deletion (task_id / node_id become NULL when the task or
|
// 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).
|
// exploration node is deleted). taskID and nodeID may be 0 (stored as NULL).
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+135
-1
@@ -1,6 +1,9 @@
|
|||||||
package db
|
package db
|
||||||
|
|
||||||
import "testing"
|
import (
|
||||||
|
"fmt"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
func TestExplorationFlow(t *testing.T) {
|
func TestExplorationFlow(t *testing.T) {
|
||||||
d, err := Open(testDSN(t))
|
d, err := Open(testDSN(t))
|
||||||
@@ -112,3 +115,134 @@ func TestExplorationFlow(t *testing.T) {
|
|||||||
t.Fatalf("stats: %+v", st)
|
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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -29,8 +29,12 @@ target AS (
|
|||||||
JOIN task_scope ts ON (
|
JOIN task_scope ts ON (
|
||||||
(ts.kind='company' AND a.company_id = ts.company_id)
|
(ts.kind='company' AND a.company_id = ts.company_id)
|
||||||
OR (ts.kind='root_domain' AND a.root_domain = ts.domain)
|
OR (ts.kind='root_domain' AND a.root_domain = ts.domain)
|
||||||
OR (ts.kind='subdomain' AND a.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 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
|
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||||
UNION
|
UNION
|
||||||
@@ -53,7 +57,7 @@ tested AS (
|
|||||||
func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) {
|
func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) {
|
||||||
rows, err := s.db.Query(`WITH `+directTaskContextCTE+`
|
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,''),
|
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
|
FROM task_scope ts
|
||||||
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||||
LEFT JOIN companies c ON c.id=ts.company_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() {
|
for rows.Next() {
|
||||||
var scope TaskScope
|
var scope TaskScope
|
||||||
var companyID int64
|
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
|
return nil, err
|
||||||
}
|
}
|
||||||
if companyID > 0 {
|
if companyID > 0 {
|
||||||
@@ -175,8 +179,12 @@ context_assets AS (
|
|||||||
JOIN task_scope ts ON (
|
JOIN task_scope ts ON (
|
||||||
(ts.kind='company' AND a.company_id=ts.company_id)
|
(ts.kind='company' AND a.company_id=ts.company_id)
|
||||||
OR (ts.kind='root_domain' AND a.root_domain=ts.domain)
|
OR (ts.kind='root_domain' AND a.root_domain=ts.domain)
|
||||||
OR (ts.kind='subdomain' AND a.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 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
|
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+281
-45
@@ -1,7 +1,11 @@
|
|||||||
package db
|
package db
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"database/sql"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strconv"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -9,15 +13,29 @@ import (
|
|||||||
type Task struct {
|
type Task struct {
|
||||||
ID int64 `json:"id"`
|
ID int64 `json:"id"`
|
||||||
Name string `json:"name"` // 可选任务名称;空=未命名
|
Name string `json:"name"` // 可选任务名称;空=未命名
|
||||||
|
CategoryID *int64 `json:"category_id,omitempty"`
|
||||||
|
CategoryName string `json:"category_name,omitempty"`
|
||||||
Description string `json:"description"`
|
Description string `json:"description"`
|
||||||
Goal string `json:"goal"`
|
Goal string `json:"goal"`
|
||||||
ExplorationID int64 `json:"exploration_id"`
|
ExplorationID int64 `json:"exploration_id"`
|
||||||
Status string `json:"status"`
|
Status string `json:"status"`
|
||||||
Paused bool `json:"paused"`
|
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"`
|
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
|
||||||
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
|
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
|
||||||
CreatedAt time.Time `json:"created_at"`
|
// for ActiveLLMProfileID while older API clients still send one profile id.
|
||||||
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
|
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
|
||||||
|
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
|
||||||
|
LLMChainRevision int64 `json:"-"`
|
||||||
|
LLMFailoverState string `json:"llm_failover_state,omitempty"`
|
||||||
|
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||||
|
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
|
||||||
|
CompanyIDs []int64 `json:"company_ids,omitempty"`
|
||||||
|
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)。
|
// 任务级超时(见 docs/任务级超时与收尾设计.md)。
|
||||||
TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时
|
TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时
|
||||||
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行
|
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行
|
||||||
@@ -32,6 +50,23 @@ type Task struct {
|
|||||||
CoverageEnabled bool `json:"coverage_enabled"`
|
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.
|
// IsTerminal reports whether a task status is a terminal (finished) state.
|
||||||
// 单一真源,替换散落各处的 done/failed 硬编码判定。
|
// 单一真源,替换散落各处的 done/failed 硬编码判定。
|
||||||
func IsTerminal(status string) bool {
|
func IsTerminal(status string) bool {
|
||||||
@@ -45,6 +80,46 @@ func IsTerminal(status string) bool {
|
|||||||
// 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min,防止把 planner 打爆。
|
// 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min,防止把 planner 打爆。
|
||||||
const MinPlanHeartbeatSeconds = 600
|
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 {
|
func normalizeHeartbeat(sec int) int {
|
||||||
if sec < MinPlanHeartbeatSeconds {
|
if sec < MinPlanHeartbeatSeconds {
|
||||||
return 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) {
|
func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) {
|
||||||
<<<<<<< Updated upstream
|
|
||||||
=======
|
|
||||||
var ids []int64
|
var ids []int64
|
||||||
if llmProfileID != nil {
|
if llmProfileID != nil {
|
||||||
ids = []int64{*llmProfileID}
|
ids = []int64{*llmProfileID}
|
||||||
@@ -68,6 +141,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe
|
|||||||
// with the task/exploration row.
|
// with the task/exploration row.
|
||||||
type TaskCreateOptions struct {
|
type TaskCreateOptions struct {
|
||||||
Name string // 可选任务名称;空=未命名
|
Name string // 可选任务名称;空=未命名
|
||||||
|
CategoryID *int64
|
||||||
SourceTaskIDs []int64
|
SourceTaskIDs []int64
|
||||||
CompanyIDs []int64
|
CompanyIDs []int64
|
||||||
LLMProfileIDs []int64
|
LLMProfileIDs []int64
|
||||||
@@ -89,7 +163,6 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
opts.CompanyIDs = companyIDs
|
opts.CompanyIDs = companyIDs
|
||||||
>>>>>>> Stashed changes
|
|
||||||
tx, err := d.Begin()
|
tx, err := d.Begin()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
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 {
|
VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
<<<<<<< Updated upstream
|
|
||||||
if timeoutSeconds < 0 {
|
|
||||||
timeoutSeconds = 0
|
|
||||||
=======
|
|
||||||
if opts.TimeoutSeconds < 0 {
|
if opts.TimeoutSeconds < 0 {
|
||||||
opts.TimeoutSeconds = 0
|
opts.TimeoutSeconds = 0
|
||||||
}
|
}
|
||||||
opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds)
|
opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds)
|
||||||
coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled
|
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
|
var active *int64
|
||||||
if len(opts.LLMProfileIDs) > 0 {
|
if len(opts.LLMProfileIDs) > 0 {
|
||||||
id := opts.LLMProfileIDs[0]
|
id := opts.LLMProfileIDs[0]
|
||||||
active = &id
|
active = &id
|
||||||
}
|
}
|
||||||
t := &Task{
|
t := &Task{
|
||||||
Name: opts.Name,
|
Name: opts.Name, CategoryID: opts.CategoryID, CategoryName: categoryName,
|
||||||
Description: description, Goal: goal, ExplorationID: expID,
|
Description: description, Goal: goal, ExplorationID: expID,
|
||||||
LLMProfileID: active, ActiveLLMProfileID: active,
|
LLMProfileID: active, ActiveLLMProfileID: active,
|
||||||
LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...),
|
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...),
|
CompanyIDs: append([]int64(nil), opts.CompanyIDs...),
|
||||||
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
|
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
|
||||||
CoverageEnabled: coverageEnabled,
|
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(`
|
if err := tx.QueryRow(`
|
||||||
<<<<<<< Updated upstream
|
INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled)
|
||||||
INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6)
|
VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9)
|
||||||
RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
|
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 {
|
||||||
=======
|
|
||||||
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
|
|
||||||
return nil, err
|
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()
|
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 {
|
func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
|
||||||
if len(companyIDs) == 0 {
|
if len(companyIDs) == 0 {
|
||||||
return nil
|
return nil
|
||||||
@@ -222,12 +302,13 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
|
|||||||
return nil
|
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) {
|
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
|
||||||
var t Task
|
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 {
|
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 {
|
||||||
>>>>>>> Stashed changes
|
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
return &t, nil
|
return &t, nil
|
||||||
@@ -254,7 +335,16 @@ func (d *DB) ListTasks() ([]*Task, error) {
|
|||||||
}
|
}
|
||||||
out = append(out, t)
|
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).
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
|
if err := d.hydrateTaskContext(t); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
return t, nil
|
return t, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -275,6 +368,43 @@ func (d *DB) SetPaused(id int64, paused bool) error {
|
|||||||
return err
|
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
|
// SetStatus updates a task's lifecycle status. Entering a terminal state
|
||||||
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
|
// (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
|
// 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
|
return deadline, err
|
||||||
}
|
}
|
||||||
|
|
||||||
// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/
|
// DeleteTask preserves the historical behavior: remove the task and its
|
||||||
// activity via ON DELETE CASCADE). The task row is removed first so the
|
// exploration subgraph while retaining global assets.
|
||||||
// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete.
|
|
||||||
// Global assets are NOT touched.
|
|
||||||
func (d *DB) DeleteTask(id int64) error {
|
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()
|
tx, err := d.Begin()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return result, err
|
||||||
}
|
}
|
||||||
defer tx.Rollback()
|
defer tx.Rollback()
|
||||||
var expID int64
|
var expID int64
|
||||||
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil {
|
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 FOR UPDATE`, id).Scan(&expID); err != nil {
|
||||||
if err.Error() == "sql: no rows in result set" {
|
if err == sql.ErrNoRows {
|
||||||
return nil // already gone
|
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 {
|
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 {
|
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()
|
||||||
}
|
}
|
||||||
+506
-1
@@ -1,6 +1,12 @@
|
|||||||
package db
|
package db
|
||||||
|
|
||||||
import "testing"
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
||||||
d, err := Open(testDSN(t))
|
d, err := Open(testDSN(t))
|
||||||
@@ -27,6 +33,23 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
|||||||
if err != nil || got == nil || !got.Paused {
|
if err != nil || got == nil || !got.Paused {
|
||||||
t.Fatalf("paused not persisted: %+v err=%v", got, err)
|
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 contains it
|
||||||
list, _ := d.ListTasks()
|
list, _ := d.ListTasks()
|
||||||
@@ -58,3 +81,485 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
|||||||
t.Fatalf("exploration should be deleted, got %d", exps)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -45,9 +45,11 @@ type chatAttachment struct {
|
|||||||
// 文件先落这里,前端按返回的 abs 绝对路径写进任务描述)
|
// 文件先落这里,前端按返回的 abs 绝对路径写进任务描述)
|
||||||
func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) {
|
func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) {
|
||||||
var sub string
|
var sub string
|
||||||
|
taskScoped := false
|
||||||
switch r.URL.Query().Get("scope") {
|
switch r.URL.Query().Get("scope") {
|
||||||
case "task":
|
case "task":
|
||||||
sub = "tasks"
|
sub = "tasks"
|
||||||
|
taskScoped = true
|
||||||
case "session":
|
case "session":
|
||||||
sub = "sessions"
|
sub = "sessions"
|
||||||
case "staging":
|
case "staging":
|
||||||
@@ -61,6 +63,17 @@ func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) {
|
|||||||
writeErr(w, 400, "非法 id")
|
writeErr(w, 400, "非法 id")
|
||||||
return
|
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")
|
dir := filepath.Join(s.m.dir, sub, id, "uploads")
|
||||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
writeErr(w, 500, "建目录失败: "+err.Error())
|
writeErr(w, 500, "建目录失败: "+err.Error())
|
||||||
|
|||||||
+523
-137
@@ -71,8 +71,18 @@ func preview(s string, n int) string {
|
|||||||
const (
|
const (
|
||||||
modelErrorRetries = 2 // model_error 收场后额外重试的次数
|
modelErrorRetries = 2 // model_error 收场后额外重试的次数
|
||||||
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
|
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
|
// 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
|
// (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;
|
// 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)
|
bc *Broadcaster // live activity pub/sub (SSE)
|
||||||
|
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
planner *agent.Planner
|
planner *agent.Planner
|
||||||
worker *agent.Worker
|
worker *agent.Worker
|
||||||
started sync.Map // taskID -> bool, so Run is idempotent per task
|
started sync.Map // taskID -> bool, so Run is idempotent per task
|
||||||
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
||||||
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
|
llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls
|
||||||
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
|
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,
|
// 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
|
// 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
|
execMu sync.Mutex
|
||||||
execCancel map[string]context.CancelFunc
|
execCancel map[string]context.CancelCauseFunc
|
||||||
execCtx map[string]context.Context
|
execCtx map[string]context.Context
|
||||||
|
|
||||||
// per-work (per running intent) cancel, so the planner's kill_work can stop a
|
// Per-work control lets the planner kill a worker and lets the UI pause/cancel
|
||||||
// single worker without pausing the whole task. Keyed by globally-unique intent id.
|
// one intent without pausing the whole task. The done channel closes only after
|
||||||
workMu sync.Mutex
|
// runWorkerStep has stopped writing and committed its final state.
|
||||||
workCancel map[int64]context.CancelFunc
|
workMu sync.Mutex
|
||||||
|
work map[int64]*workExecution
|
||||||
|
|
||||||
// steerBox queues planner course-corrections for a running work (keyed by intent
|
// 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
|
// 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)
|
inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain)
|
||||||
coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重)
|
coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重)
|
||||||
|
|
||||||
// resolve, if set, returns a task's dedicated planner/worker when it pins a
|
// resolve, if set, returns a task's dedicated planner/worker. In legacy mode
|
||||||
// specific LLM profile (nil,nil = use the global active pair). Set once at startup.
|
// nil,nil falls back to the global pair; an authoritative resolver uses nil,nil
|
||||||
resolve func(t *Task) (*agent.Planner, *agent.Worker)
|
// 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).
|
// 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
|
// 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).
|
// 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.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
|
// 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
|
// planner.Plan / worker.Execute), if present. Shared by Pause and the settle
|
||||||
// sequence's hard-drain backstop.
|
// sequence's hard-drain backstop.
|
||||||
func (e *Engine) cancelExec(taskID string) {
|
func (e *Engine) cancelExec(taskID string, cause error) {
|
||||||
e.execMu.Lock()
|
e.execMu.Lock()
|
||||||
if cancel := e.execCancel[taskID]; cancel != nil {
|
if cancel := e.execCancel[taskID]; cancel != nil {
|
||||||
cancel()
|
cancel(cause)
|
||||||
}
|
}
|
||||||
e.execMu.Unlock()
|
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
|
// Resume un-pauses a task and nudges a fresh planning round. The next exec under
|
||||||
// it gets a fresh (uncancelled) context.
|
// it gets a fresh (uncancelled) context.
|
||||||
func (e *Engine) Resume(t *Task) {
|
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)
|
e.paused.Delete(t.ID)
|
||||||
t.Notify()
|
t.Notify()
|
||||||
}
|
}
|
||||||
@@ -166,14 +326,14 @@ func (e *Engine) execContextFor(parent context.Context, taskID string) context.C
|
|||||||
defer e.execMu.Unlock()
|
defer e.execMu.Unlock()
|
||||||
if e.IsPaused(taskID) {
|
if e.IsPaused(taskID) {
|
||||||
// never hand out a live context while paused (guards the claim→Execute race)
|
// never hand out a live context while paused (guards the claim→Execute race)
|
||||||
c, cancel := context.WithCancel(parent)
|
c, cancel := context.WithCancelCause(parent)
|
||||||
cancel()
|
cancel(agent.AbortPausedRaceGuard)
|
||||||
return c
|
return c
|
||||||
}
|
}
|
||||||
if c := e.execCtx[taskID]; c != nil && c.Err() == nil {
|
if c := e.execCtx[taskID]; c != nil && c.Err() == nil {
|
||||||
return c
|
return c
|
||||||
}
|
}
|
||||||
c, cancel := context.WithCancel(parent)
|
c, cancel := context.WithCancelCause(parent)
|
||||||
e.execCtx[taskID] = c
|
e.execCtx[taskID] = c
|
||||||
e.execCancel[taskID] = cancel
|
e.execCancel[taskID] = cancel
|
||||||
return c
|
return c
|
||||||
@@ -200,32 +360,131 @@ func (e *Engine) LastActivity(taskID string) int64 {
|
|||||||
return 0
|
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 (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) }
|
||||||
|
|
||||||
func NewEngine(m *Manager) *Engine {
|
func NewEngine(m *Manager) *Engine {
|
||||||
return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(),
|
return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(),
|
||||||
execCancel: map[string]context.CancelFunc{}, execCtx: map[string]context.Context{},
|
execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{},
|
||||||
workCancel: map[int64]context.CancelFunc{}, steerBox: map[int64][]string{}}
|
work: map[int64]*workExecution{}, steerBox: map[int64][]string{},
|
||||||
|
runtimes: map[string]*taskRuntime{}}
|
||||||
}
|
}
|
||||||
|
|
||||||
// registerWork records the cancel for the work currently running intentID.
|
// 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.workMu.Lock()
|
||||||
e.workCancel[intentID] = cancel
|
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan error, 1)}
|
||||||
e.workMu.Unlock()
|
e.workMu.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
// unregisterWork removes and releases the work's cancel (called when Execute returns).
|
// detachWork removes the live control handle once Execute has returned. complete
|
||||||
func (e *Engine) unregisterWork(intentID int64) {
|
// 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()
|
e.workMu.Lock()
|
||||||
if c := e.workCancel[intentID]; c != nil {
|
run := e.work[intentID]
|
||||||
delete(e.workCancel, intentID)
|
if run != nil {
|
||||||
c() // release context resources (no-op if already cancelled)
|
delete(e.work, intentID)
|
||||||
|
action = run.action
|
||||||
|
run.cancel(agent.AbortWorkFinished) // release resources (no-op if already cancelled)
|
||||||
}
|
}
|
||||||
e.workMu.Unlock()
|
e.workMu.Unlock()
|
||||||
e.steerMu.Lock()
|
e.steerMu.Lock()
|
||||||
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
|
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
|
||||||
e.steerMu.Unlock()
|
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
|
// 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("纠偏消息不能为空")
|
return fmt.Errorf("纠偏消息不能为空")
|
||||||
}
|
}
|
||||||
e.workMu.Lock()
|
e.workMu.Lock()
|
||||||
running := e.workCancel[intentID] != nil
|
running := e.work[intentID] != nil
|
||||||
e.workMu.Unlock()
|
e.workMu.Unlock()
|
||||||
if !running {
|
if !running {
|
||||||
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
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.
|
// The work's agent-core session honors ctx cancellation and aborts promptly.
|
||||||
func (e *Engine) KillWork(intentID int64) error {
|
func (e *Engine) KillWork(intentID int64) error {
|
||||||
e.workMu.Lock()
|
e.workMu.Lock()
|
||||||
c := e.workCancel[intentID]
|
run := e.work[intentID]
|
||||||
e.workMu.Unlock()
|
e.workMu.Unlock()
|
||||||
if c == nil {
|
if run == nil {
|
||||||
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
||||||
}
|
}
|
||||||
c()
|
run.cancel(agent.AbortKilledByPlanner)
|
||||||
return nil
|
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,
|
// emitActivity persists one captured step AND fans it out to live subscribers,
|
||||||
// from a single point so storage and the SSE stream never diverge.
|
// 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)
|
id, err := e.appendActivity(t, r)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// NO LONGER SILENT: dropping a record breaks command↔result pairing in the
|
// 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",
|
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)
|
t.ID, n, r.Worker, r.Kind, r.Tool, r.ToolUseID, dropReason(err), preview(r.Summary, 80), err, diag)
|
||||||
e.touch(t.ID)
|
e.touch(t.ID)
|
||||||
return
|
return r
|
||||||
}
|
}
|
||||||
r.ID = id
|
r.ID = id
|
||||||
if r.CreatedAt.IsZero() {
|
if r.CreatedAt.IsZero() {
|
||||||
@@ -348,6 +607,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
|||||||
}
|
}
|
||||||
e.bc.Publish(t.ID, r)
|
e.bc.Publish(t.ID, r)
|
||||||
e.touch(t.ID)
|
e.touch(t.ID)
|
||||||
|
return r
|
||||||
}
|
}
|
||||||
|
|
||||||
// appendActivity persists one activity row, retrying briefly on write failure.
|
// 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.
|
// 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)) {
|
func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) {
|
||||||
e.resolve = fn
|
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
|
// 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).
|
// pair. The resolver returning nil means "no override" (fall back to global).
|
||||||
func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) {
|
func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) {
|
||||||
if e.resolve != nil {
|
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
|
return p, w
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -410,6 +680,14 @@ func (e *Engine) Ready() bool {
|
|||||||
return p != nil && w != nil
|
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".
|
// Mode reports "llm" when ready, else "idle".
|
||||||
func (e *Engine) Mode() string {
|
func (e *Engine) Mode() string {
|
||||||
if e.Ready() {
|
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
|
// but no-op until an LLM is configured (so a task created while idle picks up
|
||||||
// automatically once LLM is set from the UI).
|
// automatically once LLM is set from the UI).
|
||||||
func (e *Engine) Run(ctx context.Context, t *Task) {
|
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 {
|
if _, loaded := e.started.LoadOrStore(t.ID, true); loaded {
|
||||||
|
e.deleteMu.RUnlock()
|
||||||
t.Notify() // already running — just nudge a planning round
|
t.Notify() // already running — just nudge a planning round
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
rt := e.registerTaskRoutines(ctx, t.ID, 1+workers)
|
||||||
|
e.deleteMu.RUnlock()
|
||||||
e.touch(t.ID)
|
e.touch(t.ID)
|
||||||
go e.plannerLoop(ctx, t)
|
runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) })
|
||||||
workers := e.m.Workers()
|
|
||||||
for i := 0; i < workers; i++ {
|
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;去重)
|
e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重)
|
||||||
// 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已
|
// 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已
|
||||||
@@ -491,10 +778,13 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
|||||||
if e.IsPaused(t.ID) {
|
if e.IsPaused(t.ID) {
|
||||||
return // user-paused: don't plan
|
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
|
// 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
|
// 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).
|
// re-plan (it would burn an LLM round and re-confirm a settled result).
|
||||||
if isTerminalStatus(t.Status) {
|
if isTerminalStatus(t.lifecycleSnapshot().Status) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
|
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
|
||||||
@@ -502,10 +792,17 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
|||||||
if e.isSettling(t.ID) {
|
if e.isSettling(t.ID) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if !e.beginTaskOperation(t.ID) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer e.decInflight(t.ID)
|
||||||
e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||||
e.touch(t.ID)
|
e.touch(t.ID)
|
||||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
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)
|
log.Printf("[planner] task %s 规划中…(%s 触发)", t.ID, src)
|
||||||
// round marker: each Plan() is one planner round; emit a boundary so the
|
// round marker: each Plan() is one planner round; emit a boundary so the
|
||||||
// UI can separate rounds in the transcript (kind='round').
|
// 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
|
// what fired this round (worker done / finding; may be several — debounce
|
||||||
// coalesces a burst; empty for time/heartbeat wakes).
|
// coalesces a burst; empty for time/heartbeat wakes).
|
||||||
triggers := t.drainTriggers()
|
triggers := t.drainTriggers()
|
||||||
e.incInflight(t.ID)
|
|
||||||
taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64)
|
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)
|
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 {
|
switch {
|
||||||
case err != nil && ectx.Err() == nil:
|
case err != nil && ectx.Err() == nil:
|
||||||
log.Printf("[planner] task %s 规划出错: %v", t.ID, err)
|
log.Printf("[planner] task %s 规划出错: %v", t.ID, err)
|
||||||
@@ -530,7 +827,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
|||||||
// 任务已判完成 → 立刻取消在跑的 worker:它们手头的意图跑出来也没意义了。
|
// 任务已判完成 → 立刻取消在跑的 worker:它们手头的意图跑出来也没意义了。
|
||||||
// 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支
|
// 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支
|
||||||
// 归为 stopped(而非 blocked)。
|
// 归为 stopped(而非 blocked)。
|
||||||
e.cancelExec(t.ID)
|
e.cancelExec(t.ID, agent.AbortGoalMet)
|
||||||
default:
|
default:
|
||||||
log.Printf("[planner] task %s 规划完成", t.ID)
|
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
|
continue // user-paused: don't claim/execute intents
|
||||||
}
|
}
|
||||||
|
if e.IsDeleting(t.ID) {
|
||||||
|
return
|
||||||
|
}
|
||||||
if e.isSettling(t.ID) {
|
if e.isSettling(t.ID) {
|
||||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||||
return
|
return
|
||||||
@@ -584,105 +884,191 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) {
|
|||||||
}
|
}
|
||||||
continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier
|
continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier
|
||||||
}
|
}
|
||||||
intent := e.claimNext(t, name)
|
if !e.beginTaskOperation(t.ID) {
|
||||||
if intent == nil {
|
return
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
claimed := e.runWorkerStep(ctx, t, name, worker)
|
||||||
e.decInflight(t.ID)
|
e.decInflight(t.ID)
|
||||||
// CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's
|
if !claimed && sleepCtx(ctx, 800*time.Millisecond) {
|
||||||
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
return
|
||||||
// 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
|
|
||||||
}
|
}
|
||||||
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
|
}
|
||||||
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
|
}
|
||||||
if ectx.Err() != nil && e.isSettling(t.ID) {
|
|
||||||
_ = t.Store.SetIntentState(intent.ID, "exhausted")
|
// runWorkerStep claims and fully settles at most one intent. Its caller holds one
|
||||||
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote)
|
// task-operation admission for the whole claim/execute/state-write sequence so a
|
||||||
e.touch(t.ID)
|
// delete cannot observe quiescence between the LLM return and the final DB writes.
|
||||||
continue
|
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。意图结果已无意义,
|
return true
|
||||||
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
|
}
|
||||||
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
// per-work child context so the planner's kill_work can stop just this work.
|
||||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
workCtx, workCancel := context.WithCancelCause(ectx)
|
||||||
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
|
e.registerWork(intent.ID, workCancel)
|
||||||
e.touch(t.ID)
|
// wrap the guard hooks so steer_work can inject a mid-run course-correction
|
||||||
continue
|
// 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).
|
e.BeginLLMCall(t.ID)
|
||||||
if killed {
|
reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding)
|
||||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
e.EndLLMCall(t.ID)
|
||||||
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
|
}
|
||||||
e.touch(t.ID)
|
// Capture kill state before detachWork cancels workCtx. kill = this work's
|
||||||
t.Notify()
|
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
||||||
continue
|
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
|
||||||
}
|
// unregister would always be true (unregister cancels it) → every completed
|
||||||
if err != nil {
|
// work would be wrongly marked stopped.
|
||||||
log.Printf("[worker %s] intent %d: %v", name, intent.ID, err)
|
workCause := context.Cause(workCtx)
|
||||||
}
|
killed := workCtx.Err() != nil && ectx.Err() == nil
|
||||||
// terminal分流:撞步数上限 ≠ 完成。max_turns→exhausted(规划者据此知道这个方向
|
action, completeWork := e.detachWork(intent.ID)
|
||||||
// 试过但没真正做完、需换角度,而非当成已覆盖永久跳过);出错→blocked;正常→done。
|
// A caller may stop waiting and release its in-memory reservation before the
|
||||||
state := "done"
|
// agent honors cancellation. The named context cause remains authoritative and
|
||||||
|
// still settles the stopped run into a recoverable state.
|
||||||
|
if action == "" {
|
||||||
switch {
|
switch {
|
||||||
case err != nil:
|
case errors.Is(workCause, agent.AbortWorkPausedByUser):
|
||||||
state = "blocked"
|
action = "pause"
|
||||||
case reason == harness.ReasonMaxTurns:
|
case errors.Is(workCause, agent.AbortWorkCancelledByUser):
|
||||||
state = "exhausted"
|
action = "cancel"
|
||||||
log.Printf("[worker %s] intent %d 撞步数上限(exhausted),本次写回 %s", name, intent.ID, wrote)
|
|
||||||
case reason == harness.ReasonTimeout:
|
|
||||||
state = "exhausted"
|
|
||||||
log.Printf("[worker %s] intent %d 运行超时(exhausted),收尾后写回 %s", name, intent.ID, wrote)
|
|
||||||
}
|
}
|
||||||
_ = t.Store.SetIntentState(intent.ID, state)
|
}
|
||||||
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote)
|
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)
|
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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
+59
-11
@@ -43,7 +43,20 @@ func (e *Engine) inflightCounter(taskID string) *int64 {
|
|||||||
v, _ := e.inflight.LoadOrStore(taskID, new(int64))
|
v, _ := e.inflight.LoadOrStore(taskID, new(int64))
|
||||||
return v.(*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) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) }
|
||||||
func (e *Engine) inflightCount(taskID string) int64 {
|
func (e *Engine) inflightCount(taskID string) int64 {
|
||||||
return atomic.LoadInt64(e.inflightCounter(taskID))
|
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 {
|
if v, ok := e.deadline.Load(t.ID); ok {
|
||||||
return v.(int64)
|
return v.(int64)
|
||||||
}
|
}
|
||||||
if t.DeadlineAt > 0 {
|
deadlineAt := t.lifecycleSnapshot().DeadlineAt
|
||||||
e.deadline.Store(t.ID, t.DeadlineAt)
|
if deadlineAt > 0 {
|
||||||
return t.DeadlineAt
|
e.deadline.Store(t.ID, deadlineAt)
|
||||||
|
return deadlineAt
|
||||||
}
|
}
|
||||||
return 0
|
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)
|
// 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.
|
// of a timeout task, once per process. No-op when the task has no timeout.
|
||||||
func (e *Engine) stampFirstRun(t *Task) {
|
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 {
|
if t == nil || t.TimeoutSeconds <= 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded {
|
e.deleteMu.RLock()
|
||||||
|
if e.IsDeleting(t.ID) {
|
||||||
|
e.deleteMu.RUnlock()
|
||||||
return
|
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
|
// 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 {
|
for e.inflightCount(t.ID) > 0 {
|
||||||
if time.Now().After(hardStop) {
|
if time.Now().After(hardStop) {
|
||||||
log.Printf("[deadline] task %s drain 超时(%s),硬取消在跑 run", t.ID, settleDrainGrace)
|
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 分支一点时间落库/归类
|
_ = sleepCtx(ctx, 3*time.Second) // 给 worker 分支一点时间落库/归类
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
@@ -170,6 +204,10 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) {
|
|||||||
|
|
||||||
// ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。
|
// ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。
|
||||||
met := e.runFinalPlannerRound(ctx, t)
|
met := e.runFinalPlannerRound(ctx, t)
|
||||||
|
if !e.beginTaskOperation(t.ID) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer e.decInflight(t.ID)
|
||||||
|
|
||||||
// ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done,
|
// ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done,
|
||||||
// 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。
|
// 守卫(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
|
// 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.
|
// 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) {
|
func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) {
|
||||||
|
if e.IsDeleting(t.ID) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
planner, _ := e.snapshotFor(t)
|
planner, _ := e.snapshotFor(t)
|
||||||
for planner == nil {
|
for planner == nil {
|
||||||
if sleepCtx(ctx, deadlinePollInterval) {
|
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)) {
|
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
if e.IsDeleting(t.ID) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
planner, _ = e.snapshotFor(t)
|
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) }
|
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||||
Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))})
|
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)
|
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)
|
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 {
|
if err != nil {
|
||||||
log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err)
|
log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err)
|
||||||
} else if met {
|
} else if met {
|
||||||
|
|||||||
+330
-59
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log"
|
"log"
|
||||||
|
"sort"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
@@ -25,25 +26,324 @@ type goalSpec struct {
|
|||||||
//
|
//
|
||||||
// 异步(goroutine)所以调用方立即返回,两条路径行为一致:秒建任务、后台拆目标。
|
// 异步(goroutine)所以调用方立即返回,两条路径行为一致:秒建任务、后台拆目标。
|
||||||
func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
|
func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
|
||||||
|
if !s.engine.beginTaskOperation(t.ID) {
|
||||||
|
return
|
||||||
|
}
|
||||||
s.seed(t, seedText)
|
s.seed(t, seedText)
|
||||||
if seedFirstIntent {
|
if seedFirstIntent {
|
||||||
s.seedFirstIntent(t)
|
s.seedFirstIntent(t)
|
||||||
}
|
}
|
||||||
go func() {
|
s.engine.decInflight(t.ID)
|
||||||
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
if _, err := s.admitTask(t, "bootstrap"); err != nil {
|
||||||
Summary: "第 0 轮目标拆解"})
|
log.Printf("[concurrency] task %s 启动失败: %v", t.ID, err)
|
||||||
goals := s.createGoals(s.ctx, t, func(r db.Activity) {
|
}
|
||||||
s.engine.emitActivity(t, r)
|
}
|
||||||
})
|
|
||||||
for _, g := range goals {
|
func (s *Server) startTaskEngine(t *Task) {
|
||||||
summary := g.Text
|
ctx := s.engine.execContextFor(s.ctx, t.ID)
|
||||||
if g.VulnClass != "" {
|
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
||||||
summary = fmt.Sprintf("[%s] %s", g.VulnClass, g.Text)
|
return
|
||||||
}
|
}
|
||||||
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary})
|
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、
|
// reviveTask 让一个已停下的任务重新跑起来:把终态(done/failed/timeout)拉回 running、
|
||||||
@@ -56,22 +356,8 @@ func (s *Server) reviveTask(t *Task) {
|
|||||||
if t == nil {
|
if t == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
// 终态 → 拉回 running(SetTaskStatus 会清 completed_at 并同步内存 handle 的 Status,
|
if _, err := s.admitTask(t, "resume"); err != nil {
|
||||||
// 使 planner/worker 循环的终态门放行)。
|
log.Printf("[revive] task %s 恢复失败: %v", t.ID, err)
|
||||||
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)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -83,24 +369,25 @@ func (s *Server) reviveTask(t *Task) {
|
|||||||
// Returns the seeded specs so callers can emit activity records for them.
|
// 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.
|
// 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 {
|
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),
|
// 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,
|
// 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
|
// so FromEnv returned empty and every task silently fell back to the crude rule
|
||||||
// splitter (which shredded URLs / made meaningless 2-way splits).
|
// splitter (which shredded URLs / made meaningless 2-way splits).
|
||||||
var specs []goalSpec
|
var specs []goalSpec
|
||||||
if cfg, ok := s.goalsLLMConfig(t); ok {
|
var as *db.AssetStore
|
||||||
var as *db.AssetStore
|
if s.m != nil {
|
||||||
if s.m != nil {
|
as = s.m.Assets()
|
||||||
as = s.m.Assets()
|
}
|
||||||
}
|
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||||
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
s.engine.BeginLLMCall(t.ID)
|
||||||
// DecomposeGoals persists each decomposed goal via the set_goals tool straight
|
decomposed := agent.DecomposeGoalsWithProvider(ctx, s.agentsForTask(t).runtime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit)
|
||||||
// into t.Store; it returns the specs it wrote so we can emit activity + spot the
|
s.engine.EndLLMCall(t.ID)
|
||||||
// empty (LLM error / no provider) case below. No write-back needed here anymore.
|
for _, g := range decomposed {
|
||||||
for _, g := range agent.DecomposeGoals(ctx, cfg, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) {
|
if strings.TrimSpace(g.Text) != "" {
|
||||||
if strings.TrimSpace(g.Text) != "" {
|
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
||||||
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(specs) == 0 {
|
if len(specs) == 0 {
|
||||||
@@ -119,19 +406,3 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity
|
|||||||
}
|
}
|
||||||
return specs
|
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()
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestCreateGoalsNilTask(t *testing.T) {
|
||||||
|
if got := (&Server{}).createGoals(context.Background(), nil, nil); len(got) != 0 {
|
||||||
|
t.Fatalf("nil task produced goals: %+v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,100 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/Autumn-27/artex/db"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestInheritedActivityDetailAndRelationDeletion(t *testing.T) {
|
||||||
|
m, err := NewManager(t.TempDir(), "")
|
||||||
|
if err != nil {
|
||||||
|
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||||
|
}
|
||||||
|
defer m.Close()
|
||||||
|
|
||||||
|
source, err := m.CreateTaskWithOptions("detail source", "source goal", db.TaskCreateOptions{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
sourceID := mustTaskID(t, source.ID)
|
||||||
|
current, err := m.CreateTaskWithOptions("detail current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{sourceID}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
_, _ = m.DeleteTask(current.ID, DeleteTaskOptions{})
|
||||||
|
_, _ = m.DeleteTask(source.ID, DeleteTaskOptions{})
|
||||||
|
})
|
||||||
|
|
||||||
|
plannerStep, err := source.Store.AppendActivity(db.Activity{
|
||||||
|
Worker: "planner", Kind: "text", Summary: "private planner", Detail: "private planner detail",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
intentID, err := source.Store.AddIntent(map[string]any{"summary": "completed source work"}, 1, nil, "planner")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := source.Store.SetIntentState(intentID, "done"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
workerStep, err := source.Store.AppendActivity(db.Activity{
|
||||||
|
NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "shared result", Detail: "shared worker detail",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||||
|
token, err := signJWT(s.jwtKey)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
fetchDetail := func(stepID int64) string {
|
||||||
|
req := httptest.NewRequest(http.MethodGet,
|
||||||
|
fmt.Sprintf("/api/exploration/activity/%d?task=%s", stepID, current.ID), nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
s.Handler().ServeHTTP(rec, req)
|
||||||
|
if rec.Code != http.StatusOK {
|
||||||
|
t.Fatalf("activity detail %d: status=%d body=%s", stepID, rec.Code, rec.Body.String())
|
||||||
|
}
|
||||||
|
var body struct {
|
||||||
|
Detail string `json:"detail"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return body.Detail
|
||||||
|
}
|
||||||
|
if got := fetchDetail(plannerStep); got != "" {
|
||||||
|
t.Fatalf("source planner transcript leaked through API: %q", got)
|
||||||
|
}
|
||||||
|
if got := fetchDetail(workerStep); got != "shared worker detail" {
|
||||||
|
t.Fatalf("terminal inherited worker detail missing: %q", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := m.DeleteTask(source.ID, DeleteTaskOptions{}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
currentAfter, ok := m.Task(current.ID)
|
||||||
|
if !ok || len(currentAfter.SourceTaskIDs) != 0 {
|
||||||
|
t.Fatalf("deleted source remained in live task DTO state: %+v", currentAfter)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func mustTaskID(t *testing.T, raw string) int64 {
|
||||||
|
t.Helper()
|
||||||
|
var id int64
|
||||||
|
if _, err := fmt.Sscan(raw, &id); err != nil || id <= 0 {
|
||||||
|
t.Fatalf("bad task id %q: %v", raw, err)
|
||||||
|
}
|
||||||
|
return id
|
||||||
|
}
|
||||||
@@ -0,0 +1,117 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Autumn-27/artex/db"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestInheritedDTOProvenance(t *testing.T) {
|
||||||
|
createdAt := time.Date(2026, time.August, 19, 12, 0, 0, 0, time.UTC)
|
||||||
|
node := &db.Node{
|
||||||
|
ID: 41,
|
||||||
|
Kind: db.KindFinding,
|
||||||
|
Payload: json.RawMessage(`{"vulnclass":"xss","severity":"high","summary":"source finding"}`),
|
||||||
|
State: "confirmed",
|
||||||
|
CreatedAt: createdAt,
|
||||||
|
SourceTaskID: 7,
|
||||||
|
Inherited: true,
|
||||||
|
}
|
||||||
|
|
||||||
|
nodeDTO := taskNodeDTO(node)
|
||||||
|
if !nodeDTO.Inherited || nodeDTO.SourceTaskID != "7" {
|
||||||
|
t.Fatalf("node provenance lost: %+v", nodeDTO)
|
||||||
|
}
|
||||||
|
finding := findingDTOsForOwner("7", "source", []*db.Node{node}, nil, nil)
|
||||||
|
if len(finding) != 1 || !finding[0].Inherited || finding[0].SourceTaskID != "7" || finding[0].TaskID != "7" {
|
||||||
|
t.Fatalf("finding provenance lost: %+v", finding)
|
||||||
|
}
|
||||||
|
|
||||||
|
intentID := int64(44)
|
||||||
|
activity := activityDTO(db.Activity{
|
||||||
|
ID: 45,
|
||||||
|
NodeID: &intentID,
|
||||||
|
CreatedAt: createdAt,
|
||||||
|
SourceTaskID: 7,
|
||||||
|
Inherited: true,
|
||||||
|
})
|
||||||
|
if !activity.Inherited || activity.SourceTaskID != "7" || activity.IntentID != "44" {
|
||||||
|
t.Fatalf("activity provenance lost: %+v", activity)
|
||||||
|
}
|
||||||
|
assetRef := coverageAssetRefDTO(db.AssetRef{
|
||||||
|
ID: 46, Kind: db.KindFact, State: "confirmed", Summary: "source fact", SourceTaskID: 7, Inherited: true,
|
||||||
|
})
|
||||||
|
if !assetRef.Inherited || assetRef.SourceTaskID != "7" {
|
||||||
|
t.Fatalf("asset ref provenance lost: %+v", assetRef)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInheritedIntentResultOnlyAllowsHistory(t *testing.T) {
|
||||||
|
for _, state := range []string{"done", "blocked", "exhausted", "stopped"} {
|
||||||
|
if !inheritedIntentResult(state) {
|
||||||
|
t.Errorf("terminal state %q should be inherited history", state)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, state := range []string{"open", "running", ""} {
|
||||||
|
if inheritedIntentResult(state) {
|
||||||
|
t.Errorf("live state %q must not enter inherited history", state)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInheritedGraphSnapshotFiltersLiveIntentsAndEdges(t *testing.T) {
|
||||||
|
nodes := []*db.Node{
|
||||||
|
{ID: 1, Kind: db.KindFact, State: "confirmed"},
|
||||||
|
{ID: 2, Kind: db.KindIntent, State: "running"},
|
||||||
|
{ID: 3, Kind: db.KindIntent, State: "done"},
|
||||||
|
{ID: 4, Kind: db.KindFinding, State: "confirmed"},
|
||||||
|
}
|
||||||
|
edges := []db.Edge{
|
||||||
|
{From: 1, Rel: db.RelDerivedFrom, To: 2},
|
||||||
|
{From: 1, Rel: db.RelDerivedFrom, To: 3},
|
||||||
|
{From: 3, Rel: db.RelYields, To: 4},
|
||||||
|
}
|
||||||
|
|
||||||
|
gotNodes, gotEdges := inheritedGraphSnapshot(nodes, edges, 77)
|
||||||
|
if len(gotNodes) != 3 {
|
||||||
|
t.Fatalf("expected fact, terminal intent and finding; got %+v", gotNodes)
|
||||||
|
}
|
||||||
|
for _, node := range gotNodes {
|
||||||
|
if node.ID == 2 {
|
||||||
|
t.Fatal("running inherited intent leaked into graph snapshot")
|
||||||
|
}
|
||||||
|
if !node.Inherited || node.SourceTaskID != 77 {
|
||||||
|
t.Fatalf("source provenance missing: %+v", node)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(gotEdges) != 2 {
|
||||||
|
t.Fatalf("edge referencing hidden live intent was not removed: %+v", gotEdges)
|
||||||
|
}
|
||||||
|
for _, edge := range gotEdges {
|
||||||
|
if edge.From == 2 || edge.To == 2 {
|
||||||
|
t.Fatalf("live intent id leaked through edge: %+v", edge)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFindingProvenanceInTask(t *testing.T) {
|
||||||
|
task := &Task{ID: "10", SourceTaskIDs: []int64{7, 8}}
|
||||||
|
|
||||||
|
local := int64(10)
|
||||||
|
if source, inherited, allowed := findingProvenanceInTask(task, &local); !allowed || inherited || source != "" {
|
||||||
|
t.Fatalf("local finding: source=%q inherited=%v allowed=%v", source, inherited, allowed)
|
||||||
|
}
|
||||||
|
sourceID := int64(7)
|
||||||
|
if source, inherited, allowed := findingProvenanceInTask(task, &sourceID); !allowed || !inherited || source != "7" {
|
||||||
|
t.Fatalf("direct source finding: source=%q inherited=%v allowed=%v", source, inherited, allowed)
|
||||||
|
}
|
||||||
|
unrelated := int64(6)
|
||||||
|
if _, _, allowed := findingProvenanceInTask(task, &unrelated); allowed {
|
||||||
|
t.Fatal("unrelated finding must not be readable through task context")
|
||||||
|
}
|
||||||
|
if _, _, allowed := findingProvenanceInTask(task, nil); allowed {
|
||||||
|
t.Fatal("detached finding must not be readable through task context")
|
||||||
|
}
|
||||||
|
}
|
||||||
+850
-63
File diff suppressed because it is too large.
Load diff
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+14
-9
@@ -199,12 +199,13 @@ func (s *Server) toolListTasks() actool.CoreTool {
|
|||||||
if t.ParentRef != "" {
|
if t.ParentRef != "" {
|
||||||
row["parent_ref"] = t.ParentRef
|
row["parent_ref"] = t.ParentRef
|
||||||
}
|
}
|
||||||
if t.LLMProfileID == nil {
|
llmState := t.llmStateSnapshot()
|
||||||
|
if llmState.ProfileID == nil {
|
||||||
row["llm_profile"] = "(激活配置)"
|
row["llm_profile"] = "(激活配置)"
|
||||||
} else if n, ok := profName[*t.LLMProfileID]; ok {
|
} else if n, ok := profName[*llmState.ProfileID]; ok {
|
||||||
row["llm_profile"] = n
|
row["llm_profile"] = n
|
||||||
} else {
|
} else {
|
||||||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID)
|
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *llmState.ProfileID)
|
||||||
}
|
}
|
||||||
out = append(out, row)
|
out = append(out, row)
|
||||||
}
|
}
|
||||||
@@ -301,10 +302,13 @@ func (s *Server) toolPauseTask() actool.CoreTool {
|
|||||||
TaskID string `json:"task_id"`
|
TaskID string `json:"task_id"`
|
||||||
}
|
}
|
||||||
_ = json.Unmarshal(in, &a)
|
_ = 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
|
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
|
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.
|
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
|
||||||
func (s *Server) deriveTaskStatus(t *Task) string {
|
func (s *Server) deriveTaskStatus(t *Task) string {
|
||||||
|
lifecycle := t.lifecycleSnapshot()
|
||||||
switch {
|
switch {
|
||||||
case isTerminalStatus(t.Status):
|
case isTerminalStatus(lifecycle.Status):
|
||||||
return t.Status
|
return lifecycle.Status
|
||||||
case t.Paused || s.engine.IsPaused(t.ID):
|
case lifecycle.Paused || s.engine.IsPaused(t.ID):
|
||||||
return "paused"
|
return "paused"
|
||||||
case s.engine.Ready() && s.engine.Started(t.ID):
|
case s.engine.ReadyFor(t) && s.engine.Started(t.ID):
|
||||||
return "running"
|
return "running"
|
||||||
}
|
}
|
||||||
return "created"
|
return "created"
|
||||||
|
|||||||
+24
-15
@@ -50,6 +50,7 @@ func (sc *Scheduler) Run(ctx context.Context) {
|
|||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
return
|
return
|
||||||
case <-t.C:
|
case <-t.C:
|
||||||
|
sc.s.reconcileConcurrency() // 并发上限:有空位就把排队任务补位启动
|
||||||
sc.step()
|
sc.step()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -155,19 +156,22 @@ func (sc *Scheduler) fireFindings(triggers []*db.AgentTrigger) {
|
|||||||
want = append(want, tr)
|
want = append(want, tr)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(want) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64)
|
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64)
|
||||||
events, err := sc.pg.NewFindingsSince(last)
|
events, err := sc.pg.NewFindingsSince(last)
|
||||||
if err != nil || len(events) == 0 {
|
if err != nil || len(events) == 0 {
|
||||||
return
|
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
|
maxID := last
|
||||||
for _, e := range events {
|
for _, e := range events {
|
||||||
if e.NodeID > maxID {
|
if e.NodeID > maxID {
|
||||||
maxID = e.NodeID
|
maxID = e.NodeID
|
||||||
}
|
}
|
||||||
|
if len(want) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
msgCtx := fmt.Sprintf("\n\n【本次由任务发现 finding 触发】\n任务: #%d %s(目标:%s)\n发现: [%s/%s] %s",
|
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)
|
e.TaskID, e.TaskDesc, e.TaskGoal, e.VulnClass, e.Severity, e.Summary)
|
||||||
for _, tr := range want {
|
for _, tr := range want {
|
||||||
@@ -185,13 +189,12 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) {
|
|||||||
want = append(want, tr)
|
want = append(want, tr)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(want) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
events, err := sc.pg.MetGoals()
|
events, err := sc.pg.MetGoals()
|
||||||
if err != nil || len(events) == 0 {
|
if err != nil || len(events) == 0 {
|
||||||
return
|
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()
|
fired := sc.firedGoalSet()
|
||||||
changed := false
|
changed := false
|
||||||
for _, e := range events {
|
for _, e := range events {
|
||||||
@@ -200,6 +203,9 @@ func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) {
|
|||||||
}
|
}
|
||||||
fired[e.NodeID] = true
|
fired[e.NodeID] = true
|
||||||
changed = true
|
changed = true
|
||||||
|
if len(want) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s(目标:%s)\n达成目标: %s",
|
msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s(目标:%s)\n达成目标: %s",
|
||||||
e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary)
|
e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary)
|
||||||
for _, tr := range want {
|
for _, tr := range want {
|
||||||
@@ -220,19 +226,20 @@ func (sc *Scheduler) fireTaskTimeouts(triggers []*db.AgentTrigger) {
|
|||||||
want = append(want, tr)
|
want = append(want, tr)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(want) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64)
|
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64)
|
||||||
events, err := sc.pg.TimedOutTasksSince(last)
|
events, err := sc.pg.TimedOutTasksSince(last)
|
||||||
if err != nil || len(events) == 0 {
|
if err != nil || len(events) == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// Advance the watermark even with no active trigger — see fireFindings.
|
||||||
maxID := last
|
maxID := last
|
||||||
for _, e := range events {
|
for _, e := range events {
|
||||||
if e.NodeID > maxID {
|
if e.NodeID > maxID {
|
||||||
maxID = e.NodeID
|
maxID = e.NodeID
|
||||||
}
|
}
|
||||||
|
if len(want) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s(目标:%s)",
|
msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s(目标:%s)",
|
||||||
e.TaskID, e.TaskDesc, e.TaskGoal)
|
e.TaskID, e.TaskDesc, e.TaskGoal)
|
||||||
for _, tr := range want {
|
for _, tr := range want {
|
||||||
@@ -251,19 +258,20 @@ func (sc *Scheduler) fireTaskCreates(triggers []*db.AgentTrigger) {
|
|||||||
want = append(want, tr)
|
want = append(want, tr)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(want) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTaskCreate), 10, 64)
|
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTaskCreate), 10, 64)
|
||||||
events, err := sc.pg.NewTasksSince(last)
|
events, err := sc.pg.NewTasksSince(last)
|
||||||
if err != nil || len(events) == 0 {
|
if err != nil || len(events) == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// Advance the watermark even with no active trigger — see fireFindings.
|
||||||
maxID := last
|
maxID := last
|
||||||
for _, e := range events {
|
for _, e := range events {
|
||||||
if e.NodeID > maxID {
|
if e.NodeID > maxID {
|
||||||
maxID = e.NodeID
|
maxID = e.NodeID
|
||||||
}
|
}
|
||||||
|
if len(want) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
msgCtx := fmt.Sprintf("\n\n【本次由任务创建触发】\n任务ID: #%d\n任务描述: %s\n任务目标: %s",
|
msgCtx := fmt.Sprintf("\n\n【本次由任务创建触发】\n任务ID: #%d\n任务描述: %s\n任务目标: %s",
|
||||||
e.TaskID, e.TaskDesc, e.TaskGoal)
|
e.TaskID, e.TaskDesc, e.TaskGoal)
|
||||||
for _, tr := range want {
|
for _, tr := range want {
|
||||||
@@ -283,19 +291,20 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) {
|
|||||||
want = append(want, tr)
|
want = append(want, tr)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if len(want) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastToolCall), 10, 64)
|
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastToolCall), 10, 64)
|
||||||
events, err := sc.pg.NewToolCallsSince(last)
|
events, err := sc.pg.NewToolCallsSince(last)
|
||||||
if err != nil || len(events) == 0 {
|
if err != nil || len(events) == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// Advance the watermark even with no active trigger — see fireFindings.
|
||||||
maxID := last
|
maxID := last
|
||||||
for _, e := range events {
|
for _, e := range events {
|
||||||
if e.NodeID > maxID {
|
if e.NodeID > maxID {
|
||||||
maxID = e.NodeID
|
maxID = e.NodeID
|
||||||
}
|
}
|
||||||
|
if len(want) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
errTag := ""
|
errTag := ""
|
||||||
if e.ToolIsErr {
|
if e.ToolIsErr {
|
||||||
errTag = "[error] "
|
errTag = "[error] "
|
||||||
|
|||||||
@@ -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())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+23
-42
@@ -11,6 +11,7 @@ import (
|
|||||||
"database/sql"
|
"database/sql"
|
||||||
"encoding/hex"
|
"encoding/hex"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"io/fs"
|
"io/fs"
|
||||||
@@ -28,9 +29,9 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"github.com/Autumn-27/artex/db"
|
||||||
"github.com/Autumn-27/norma/permission"
|
"github.com/Autumn-27/norma/permission"
|
||||||
actool "github.com/Autumn-27/norma/tool"
|
actool "github.com/Autumn-27/norma/tool"
|
||||||
"github.com/Autumn-27/artex/db"
|
|
||||||
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
|
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
|
||||||
_ "modernc.org/sqlite"
|
_ "modernc.org/sqlite"
|
||||||
)
|
)
|
||||||
@@ -307,15 +308,6 @@ VALUES(?,?,?,?,?,?,?)`,
|
|||||||
return
|
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} {
|
for _, h := range []string{reqB.hash, respB.hash} {
|
||||||
if h == "" {
|
if h == "" {
|
||||||
continue
|
continue
|
||||||
@@ -325,7 +317,6 @@ VALUES(?,?,?,?,?,?,?)`,
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
>>>>>>> Stashed changes
|
|
||||||
|
|
||||||
if t.fts {
|
if t.fts {
|
||||||
// The index is fed from memory, so a body that spilled to _blobs is still
|
// 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()
|
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 {
|
func sanitize(s string) string {
|
||||||
r := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "?", "_", "*", "_", "\"", "_", "<", "_", ">", "_", "|", "_")
|
r := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "?", "_", "*", "_", "\"", "_", "<", "_", ">", "_", "|", "_")
|
||||||
out := r.Replace(s)
|
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
|
// so picking "api.example.com" never sweeps "api.example.com.cn". Returns rows
|
||||||
// deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass).
|
// deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass).
|
||||||
func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) {
|
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)
|
stage, err := t.StageDeleteHostsExact(hosts)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
@@ -1028,30 +1031,9 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error
|
|||||||
stage.tx = tx
|
stage.tx = tx
|
||||||
for _, h := range unique {
|
for _, h := range unique {
|
||||||
n, err := t.deleteWhere(tx, `host=?`, h)
|
n, err := t.deleteWhere(tx, `host=?`, h)
|
||||||
>>>>>>> Stashed changes
|
|
||||||
if err != nil {
|
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:<hex>" 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
|
stage.deleted += n
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1122,7 +1104,6 @@ func (s *HostDeleteStage) Commit() error {
|
|||||||
// left the buckets behind permanently. Best-effort: walk errors only skip the
|
// left the buckets behind permanently. Best-effort: walk errors only skip the
|
||||||
// affected entries. Callers must hold wmu so this never races a concurrent
|
// affected entries. Callers must hold wmu so this never races a concurrent
|
||||||
// record() writing a fresh blob + reference.
|
// record() writing a fresh blob + reference.
|
||||||
>>>>>>> Stashed changes
|
|
||||||
func (t *Traffic) gcBlobs() error {
|
func (t *Traffic) gcBlobs() error {
|
||||||
refs := make(map[string]struct{})
|
refs := make(map[string]struct{})
|
||||||
rows, err := t.db.Query(`SELECT DISTINCT hash FROM blob_refs`)
|
rows, err := t.db.Query(`SELECT DISTINCT hash FROM blob_refs`)
|
||||||
|
|||||||
@@ -2,12 +2,30 @@ package traffic
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"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
|
// TestDeleteHost verifies the delete contract: rows for hosts containing the
|
||||||
// substring are removed together with their file trees, non-matching hosts are
|
// substring are removed together with their file trees, non-matching hosts are
|
||||||
// untouched, and the count is right.
|
// 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
|
// TestDeleteHostGCBlobs verifies blob garbage collection: after a host's trees
|
||||||
// are removed, blobs referenced by no remaining exchange are deleted, while
|
// are removed, blobs referenced by no remaining exchange are deleted, while
|
||||||
// blobs still referenced (including shared ones) survive.
|
// blobs still referenced (including shared ones) survive.
|
||||||
|
|||||||
Reference in new issue
Block a user