From 9793b522d93e14180dd52a90df21682f0652362d Mon Sep 17 00:00:00 2001 From: pphua Date: Thu, 20 Aug 2026 18:19:10 +0800 Subject: [PATCH] feat(tasks): add inherited context and resilient LLM execution Add read-only task context inheritance, ordered task-level LLM failover, live Worker controls, cascading task cleanup, FIFO task concurrency, cancellation diagnostics, UI updates, tests, build script, and changelog. --- CHANGELOG.md | 37 + agent/blackboard_inheritance_test.go | 198 +++++ agent/cancelcause.go | 124 ++- agent/capture.go | 13 +- agent/goals.go | 10 + agent/mainagent.go | 22 +- agent/planner.go | 21 +- agent/provider.go | 107 ++- agent/provider_quota_test.go | 131 +++ agent/terminalreason.go | 135 +-- agent/terminalreason_test.go | 161 ++-- agent/tools.go | 697 ++++++++++++---- agent/tools_insert.go | 9 +- agent/tools_overview_test.go | 31 + agent/worker.go | 14 +- build.sh | 65 ++ cmd/artex/main.go | 23 +- cmd/artex/main_test.go | 19 + db/assets.go | 196 +++-- db/assets_test.go | 10 +- db/config.go | 281 ++++++- db/config_test.go | 14 + db/conversation.go | 51 +- db/exploration.go | 375 ++++++++- db/exploration_sources.go | 350 ++++++++ db/exploration_sources_test.go | 431 ++++++++++ db/exploration_test.go | 136 ++- db/findings.go | 3 +- db/schema.sql | 80 +- db/task_assets_context.go | 215 +++++ db/task_context.go | 388 +++++++++ db/task_context_lock_test.go | 527 ++++++++++++ db/task_context_unit_test.go | 27 + db/task_delete_concurrency_test.go | 210 +++++ db/task_scope.go | 14 +- db/task_source_limit_test.go | 20 + db/tasks.go | 296 ++++++- db/tasks_test.go | 410 +++++++++- llmrec/llmrec.go | 41 +- llmrec/llmrec_test.go | 29 + server/assets.go | 7 +- server/chatupload.go | 13 + server/core_test.go | 78 +- server/dto.go | 162 ++-- server/engine.go | 564 +++++++++---- server/engine_cancelcause_test.go | 95 ++- server/engine_llm_calls_test.go | 43 + server/engine_timeout.go | 49 +- server/goals.go | 150 ++-- server/goals_test.go | 12 + server/inheritance_api_test.go | 100 +++ server/inheritance_dto_test.go | 117 +++ server/manager.go | 394 ++++++++- server/manager_delete_files_test.go | 280 +++++++ server/orchestration.go | 10 +- server/server.go | 774 +++++++++++++++--- server/server_mgmt.go | 182 ++-- server/task_delete_barrier_test.go | 220 +++++ server/task_llm.go | 462 +++++++++++ server/task_llm_test.go | 507 ++++++++++++ traffic/traffic.go | 214 ++++- traffic/traffic_test.go | 146 ++++ .../app/(main)/_components/main-content.tsx | 4 +- .../(main)/function/findings/detail/page.tsx | 147 ++-- .../tasks/detail/_tabs/assets-tab.tsx | 10 +- .../tasks/detail/_tabs/coverage-graph-tab.tsx | 21 +- .../tasks/detail/_tabs/findings-tab.tsx | 100 +-- .../tasks/detail/_tabs/report-tab.tsx | 46 +- .../tasks/detail/_tabs/sessions-tab.tsx | 422 ++++++++-- .../app/(main)/function/tasks/detail/page.tsx | 282 +++++-- web/src/app/(main)/function/tasks/page.tsx | 644 ++++++++++----- web/src/components/agent-editor.tsx | 7 +- web/src/components/exploration-graph.tsx | 21 +- web/src/components/task-llm-profile-chain.tsx | 223 +++++ web/src/components/transcript.tsx | 7 +- web/src/components/ui/checkbox.tsx | 8 +- web/src/components/ui/combobox.tsx | 7 +- web/src/lib/api.ts | 74 +- web/src/lib/mock/data.ts | 34 +- web/src/lib/mock/handler.ts | 189 ++++- web/src/lib/status.ts | 1 + web/src/lib/types.ts | 59 +- 82 files changed, 11094 insertions(+), 1712 deletions(-) create mode 100644 CHANGELOG.md create mode 100644 agent/blackboard_inheritance_test.go create mode 100644 agent/provider_quota_test.go create mode 100644 agent/tools_overview_test.go create mode 100755 build.sh create mode 100644 cmd/artex/main_test.go create mode 100644 db/exploration_sources.go create mode 100644 db/exploration_sources_test.go create mode 100644 db/task_assets_context.go create mode 100644 db/task_context.go create mode 100644 db/task_context_lock_test.go create mode 100644 db/task_context_unit_test.go create mode 100644 db/task_delete_concurrency_test.go create mode 100644 db/task_source_limit_test.go create mode 100644 llmrec/llmrec_test.go create mode 100644 server/engine_llm_calls_test.go create mode 100644 server/goals_test.go create mode 100644 server/inheritance_api_test.go create mode 100644 server/inheritance_dto_test.go create mode 100644 server/manager_delete_files_test.go create mode 100644 server/task_delete_barrier_test.go create mode 100644 server/task_llm.go create mode 100644 server/task_llm_test.go create mode 100644 web/src/components/task-llm-profile-chain.tsx diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..4c53e7e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,37 @@ +# 更新日志 + +本项目的重要变更记录在此文件中，格式参考 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/)。 + +## [Unreleased] - 2026-08-20 + +### Added + +- 新建任务支持关联多个已有任务，实时、只读继承直接来源任务的事实、漏洞、已完成意图、资产范围和黑板上下文；新任务仍使用独立的执行队列、工作目录和会话历史。 +- 任务支持有序 LLM 配置链。明确识别到供应商额度不足时自动切换到下一配置，并持久化当前配置、耗尽状态和结构化审计活动。 +- 运行中和暂停任务支持编辑 LLM 配置链、调整顺序及手动切换当前配置；自动切换、手动切换和全链耗尽会在任务详情页提示。 +- 运行中的 Worker 支持暂停、恢复和取消。暂停保留黑板数据，取消会在 Worker 停止写入后事务性清理该意图及其直接产出的事实、漏洞和执行记录。 +- 删除任务可选择同时清理关联资产、流量、漏洞和任务落地文件，并增加删除屏障、并发保护和可审计的删除统计。 +- 增加可选的任务并发上限；达到上限的新任务按 FIFO 排队，并在运行槽位释放后自动启动。 +- 增加任务资产分页、公司资产分页、任务分类统计及分页意图 Mock 契约。 +- 增加 `build.sh` 单二进制构建脚本，支持静态前端导出、资源嵌入、跨平台目标和构建版本注入。 + +### Changed + +- 任务级显式 LLM 配置链与现有全局 LLM Pool 并存；显式链继续采用严格的额度故障转移语义，无显式链时沿用 Agent 绑定和全局配置规则。 +- 主 Agent 输入框支持自动增长的多行输入，`Enter` 发送、`Shift+Enter` 换行，并避免中文输入法组合态误发送。 +- Planner、Worker 和主 Agent 共享任务级 LLM runtime；任务链使用候选模型中最小上下文窗口作为安全压缩阈值。 +- 任务详情页按真实 LLM 调用状态展示运行、空闲和暂停状态；来源任务的事实、漏洞、意图、资产引用和图谱节点统一标记来源并保持只读。 +- Agent 预览、任务报告和相关详情视图统一使用共享 Markdown 渲染组件。 +- LLM 模型配置采用卡片和抽屉交互，模型列表在抽屉内支持独立滚动。 + +### Fixed + +- 修复暂停任务后主 Agent 会话仍继续运行，以及编排 Agent 暂停任务时未同步停止主 Agent 的问题。 +- 修复任务完成、Planner、Worker 或主 Agent 实际运行时，任务状态徽标与操作按钮状态不同步的问题。 +- 修复 Worker 取消、任务删除和并发写入之间可能产生迟到黑板写入或残留文件的问题。 +- 修复 Markdown 预览失效、删除确认长文本溢出、移动端宽度和部分任务详情按钮未对齐的问题。 +- 为所有 Agent 取消路径增加具名终止原因，活动详情可显示取消方、终态、运行轮数、耗时、Token 使用量和未返回的工具调用。 +- 修复后端关闭时父 context 可能抢先覆盖具名 `shutdown` 原因的竞态，以及中文活动摘要按字节截断导致乱码的问题。 +- 修复带部分流式输出的取消事件丢失真实终止原因的问题，同时保留取消前已生成的内容。 + +[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.1...HEAD diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go new file mode 100644 index 0000000..978c8bf --- /dev/null +++ b/agent/blackboard_inheritance_test.go @@ -0,0 +1,198 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "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 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]) + } + if facts, ok := related[0]["recent_facts"].([]map[string]any); !ok || len(facts) == 0 || facts[0]["id"] != sourceFact { + t.Fatalf("related fact summary: %#v", related[0]["recent_facts"]) + } + if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding { + t.Fatalf("related finding summary: %#v", related[0]["recent_findings"]) + } + results, ok := related[0]["recent_intent_results"].([]map[string]any) + if !ok || len(results) == 0 || results[0]["id"] != sourceIntent { + t.Fatalf("terminal source intent was starved by newer live work: %#v", related[0]["recent_intent_results"]) + } + if results[0]["result_summary"] != "shared trace marker" { + t.Fatalf("terminal source intent result was not distilled: %#v", results[0]) + } + coverage, ok := overview["coverage"].(map[string]any) + if !ok { + t.Fatalf("coverage missing from overview: %#v", overview["coverage"]) + } + hosts, ok := coverage["hosts"].([]string) + if !ok || len(hosts) != 1 || hosts[0] != host || coverage["host_count"] != 1 { + t.Fatalf("inherited host context missing: %#v", coverage) + } + + facts := callReadJSON(t, tools.listFacts(), `{}`).([]any) + seenSource, seenGrand := false, false + for _, raw := range facts { + item := raw.(map[string]any) + id := int64(item["id"].(float64)) + if id == sourceFact { + seenSource = item["inherited"] == true && int64(item["source_task_id"].(float64)) == source.ID + } + seenGrand = seenGrand || id == grandFact + } + if !seenSource || seenGrand { + t.Fatalf("list_facts direct-only: source=%v grand=%v payload=%#v", seenSource, seenGrand, facts) + } + + findings := callReadJSON(t, tools.listFindings(), `{}`).([]any) + if len(findings) != 1 { + t.Fatalf("list_findings: %#v", findings) + } + finding := findings[0].(map[string]any) + if finding["inherited"] != true || int64(finding["task_id"].(float64)) != source.ID || int64(finding["intent_id"].(float64)) != sourceIntent { + t.Fatalf("inherited finding provenance: %#v", finding) + } + + detailInput, _ := json.Marshal(map[string]any{"id": sourceFact}) + detail := callReadJSON(t, tools.nodeDetail(), string(detailInput)).(map[string]any) + if detail["inherited"] != true || int64(detail["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited node detail: %#v", detail) + } + + traceInput, _ := json.Marshal(map[string]any{"intent_id": sourceIntent}) + trace := callReadJSON(t, tools.getWorkerTrace(), string(traceInput)).(map[string]any) + if trace["inherited"] != true || int64(trace["source_task_id"].(float64)) != source.ID { + t.Fatalf("inherited trace: %#v", trace) + } + steps := trace["steps"].([]any) + if len(steps) != 1 || int64(steps[0].(map[string]any)["step_id"].(float64)) != stepID { + t.Fatalf("inherited trace steps: %#v", steps) + } + output := callReadJSON(t, tools.getWorkerOutput(), string(traceInput)).(map[string]any) + if output["inherited"] != true || output["final_text"] != "shared trace full detail" { + t.Fatalf("inherited worker output: %#v", output) + } + + search := callReadJSON(t, tools.searchAllWorkerTraces(), `{"q":"shared trace marker"}`).(map[string]any) + hits := search["hits"].([]any) + if len(hits) != 1 || hits[0].(map[string]any)["inherited"] != true { + t.Fatalf("inherited trace search: %#v", hits) + } +} diff --git a/agent/cancelcause.go b/agent/cancelcause.go index cfda073..f4705bc 100644 --- a/agent/cancelcause.go +++ b/agent/cancelcause.go @@ -6,25 +6,12 @@ import ( "fmt" ) -// AbortCause names WHY a planner/worker run's context was cancelled. -// -// Background: when a run's ctx is cancelled mid-flight the harness finishes with -// ReasonAbortedStreaming / ReasonAbortedTools and Terminal.Err = context.Canceled. -// ARTEX cancels that ctx from a dozen different places (pause, delete, kill_work, -// goal-met, settle drain backstop, chat stop, process shutdown, hard run timeout, -// …) but by the time it reaches captureRun they all collapse into one bare -// context.Canceled — so the trace could only print a guess-list, not a diagnosis. -// -// Fix: every cancel site uses context.WithCancelCause / WithTimeoutCause and -// attaches one of the causes below; captureRun recovers it via context.Cause and -// prints exactly who cancelled the run, why, and what happens to the intent. -// -// When you add a new cancel site, add its cause HERE — a site that cancels -// without one shows up in the trace as "canceled_no_cause", by design. +// 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 // stable machine-readable id (logs / future filtering) - Short string // 一句话，进折叠摘要行 - Text string // 完整说明：谁取消的、为什么、这条意图接下来会怎样 + Code string + Short string + Text string } func (c *AbortCause) Error() string { return c.Text } @@ -33,84 +20,71 @@ func cause(code, short, text string) *AbortCause { return &AbortCause{Code: code, Short: short, Text: text} } -// Causef builds an ad-hoc cause with runtime detail (e.g. which intent, which -// deadline). Use the predeclared vars below when there is nothing dynamic to add. +// 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...)} } -// The full enumeration of cancel sites, grouped by which context gets cancelled. var ( - // ---- 任务级 exec ctx（Engine.cancelExec：取消该任务在跑的 planner + 全部 worker）---- - + // Task-level execution context. AbortPausedByUser = cause("paused_by_user", "用户暂停了任务", - "用户暂停了任务（前端「暂停」按钮 / POST /api/tasks/{id}/state action=pause）。本次运行被主动取消，该意图已退回 frontier(open)，恢复任务后会被重新领取、从头重跑") - - AbortPausedByOrchestrator = cause("paused_by_orchestrator", "编排 agent 调用 pause_task 暂停了任务", - "编排 agent 调用了 pause_task 工具暂停本任务。本次运行被主动取消，该意图已退回 frontier(open)，恢复后重跑") - + "用户通过任务控制接口（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。本次运行的结果不会再被使用") + "任务正在删除（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，收尾阶段已经写入的事实和资产会保留") - AbortPausedOnReload = cause("paused_on_reload", "后端启动时恢复了「已暂停」状态", - "后端启动时按库里持久化的「已暂停」状态恢复暂停。本次运行被取消（正常情况下此时并没有在跑的 run）") + // 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") - AbortGoalMet = cause("goal_met", "规划者判定目标已达成，任务落终态 done", - "规划者判定任务目标已达成、任务落终态 done，随即取消所有在跑的 worker——它们手头的意图跑完也没意义了。该意图标记为 stopped，不是失败") + // 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 资源。这不是运行中断；若它出现在中断消息中，说明取消与收场事件发生了竞态") - AbortSettleDrainTimeout = cause("settle_drain_timeout", "任务级超时收尾：drain 宽限用尽后强制取消", - "任务级超时收尾的硬兜底：到达任务 timeout 后等在跑 worker 优雅收尾，超过 90s drain 宽限仍未结束 → 强制取消。该意图标记为 exhausted，收尾阶段通常已把事实/资产写回") - - // ---- 单个 work 的 ctx（Engine.workCancel[intentID]）---- - - AbortKilledByPlanner = cause("killed_by_planner", "规划者用 kill_work 终止了这条意图", - "规划者调用 kill_work 主动终止了这条意图（判定方向跑偏/无意义，止损）。该意图标记为 stopped，不会被自动重领") - - AbortWorkFinished = cause("work_finished", "work 已正常结束，引擎释放其 context", - "work 已正常结束，引擎在 unregisterWork 里释放其 context 资源。这不是中断——若它出现在中断消息里，说明取消与收场事件存在竞态") - - AbortPausedRaceGuard = cause("paused_race_guard", "任务处于暂停态，引擎拒绝启动新 run", - "任务处于暂停态时请求执行上下文，引擎直接返回一个已取消的 context（守卫 claim→Execute 的竞态，避免暂停后仍启动新 run）。该意图会退回 frontier") - - // ---- 对话 / 主 Agent 的单轮 ctx（Server.chatCancel[busyKey]）---- - - AbortChatStoppedByUser = cause("chat_stopped_by_user", "用户点了「停止」中止本轮对话", - "用户点了「停止」，主动中止本轮对话运行（POST .../chat/stop 或 .../conversations/{id}/stop）。已产生的步骤都保留，可以直接发下一条消息") - - AbortChatTurnFinished = cause("chat_turn_finished", "本轮对话已正常结束，服务端释放其 context", - "本轮对话已正常结束，服务端在 defer 里释放其 context 资源。这不是中断——若它出现在中断消息里，说明取消与收场事件存在竞态") - - // ---- 进程级 ---- - - AbortShutdown = cause("shutdown", "后端进程正在关停（重启 / 更新 / Ctrl-C）", - "后端进程收到 SIGINT/SIGTERM 正在关停（重启 / 更新 / 手动 Ctrl-C）。所有在跑的 planner/worker 一并取消；重启后残留的 running 意图会被重置为 open 重跑") - - // ---- 单次 run 级（worker.Execute / planner.Plan 的硬兜底）---- - - AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底触发：某一步卡死了", - "单次运行的硬超时兜底被触发：正常情况下软墙钟预算(MaxDuration)会在回合边界优雅收尾，只有当某一步（模型请求或某个工具）自身卡死、久久不返回、导致回合边界检查根本跑不到时，才会走到「软预算 + 90s 宽限」的硬 cancel。请重点看中断前最后一个 tool_use——多半是它挂住了") + // Process-level and per-run hard backstop. + AbortShutdown = cause("shutdown", "后端进程正在关闭", + "后端进程收到 SIGINT 或 SIGTERM，正在重启、更新或关闭。所有运行中的 Agent 会被取消；重启后残留的 running 意图会重置为 open 并重新执行") + AbortRunHardTimeout = cause("run_hard_timeout", "单次运行的硬超时兜底已触发", + "单次运行超过软墙钟预算及额外宽限，说明模型请求或某个工具长时间没有返回，导致正常的回合边界收尾无法执行。请重点检查中断前最后一个未返回的工具调用") ) -// AbortReason resolves the cause attached to a cancelled run context. ok=false -// means the ctx carries no cause at all (nothing was cancelled). +// 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 } - if ac, ok := errors.AsType[*AbortCause](c); ok { + 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 + return "deadline_exceeded", "上游 context 到达 deadline", + "上游 context 到达 deadline，但设置方没有通过 WithTimeoutCause 附加具名原因: " + c.Error(), true case errors.Is(c, context.Canceled): - // a plain Canceled means some cancel site still lacks a cause — say so - // plainly instead of inventing one. - return "canceled_no_cause", "取消方未挂具名原因（漏网的取消点）", - "上游 context 被取消，但取消方没有用 context.WithCancelCause 挂上具名原因。" + - "这属于漏网的取消点——请在 agent/cancelcause.go 补一条枚举并接到该取消处", true + return "canceled_no_cause", "取消方未附加具名原因", + "上游 context 被取消，但取消方没有通过 context.WithCancelCause 附加具名原因；请在 agent/cancelcause.go 登记原因并接入该取消点", true + default: + return "other", firstLine(c.Error(), 80), c.Error(), true } - return "other", firstLine(c.Error(), 80), c.Error(), true } diff --git a/agent/capture.go b/agent/capture.go index bbede4f..ca36acc 100644 --- a/agent/capture.go +++ b/agent/capture.go @@ -6,10 +6,10 @@ import ( "time" "unicode/utf8" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/agentcore" "github.com/Autumn-27/norma/harness" "github.com/Autumn-27/norma/llm" - "github.com/Autumn-27/artex/db" ) // captureRun drives one agent turn-to-completion over Session.Prompt and emits a @@ -65,9 +65,6 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, tkind = kind tbuf.WriteString(text) } - - // breadcrumbs for the terminal message: when a run is aborted or times out, - // "which tool was still running, for how long" is the whole diagnosis. lastTool := &runTrace{startedAt: time.Now()} var finalText string @@ -75,7 +72,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, for ev, err := range s.Prompt(ctx, input) { if err != nil { flush() - if ctx.Err() != nil { // ctx cancelled = stopped by the engine, not a failure + if ctx.Err() != nil { // engine/user cancellation, not a provider failure sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool) rec(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}) return finalText, reason, ctx.Err() @@ -130,7 +127,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string, } flush() // flush any trailing thinking / non-final text sum, detail := ev.Terminal.Text, ev.Terminal.Text - if sum == "" { + if sum == "" || ev.Terminal.Reason == harness.ReasonAbortedTools || ev.Terminal.Reason == harness.ReasonAbortedStreaming { sum, detail = terminalText(ctx, ev.Terminal, lastTool) } u := ev.Terminal.Usage // cumulative token usage for this session @@ -162,9 +159,7 @@ func blocksText(blocks []llm.ContentBlock) string { return b.String() } -// firstLine returns a single-line, length-capped preview for the summary column. -// max counts RUNES, not bytes: these previews are mostly 中文, and byte-slicing -// used to cut a multi-byte rune in half and render as "恢�…". +// firstLine returns a single-line, rune-capped preview for the summary column. func firstLine(s string, max int) string { s = strings.TrimSpace(s) if before, _, found := strings.Cut(s, "\n"); found { diff --git a/agent/goals.go b/agent/goals.go index dbbae9f..3501157 100644 --- a/agent/goals.go +++ b/agent/goals.go @@ -84,6 +84,16 @@ type GoalSpec struct { // runtime). The returned specs are read back from the store so callers can emit // per-goal activity and detect the "LLM produced nothing" case for their fallback. func DecomposeGoals(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { + if prov == nil { + return nil + } + return DecomposeGoalsWithProvider(ctx, prov, dataDir, goalText, desc, as, ts, taskID, emit) +} + +// DecomposeGoalsWithProvider is the task-runtime variant used when a task has an +// ordered provider chain. It preserves the same tools and write behavior while +// letting the caller own provider selection/failover. +func DecomposeGoalsWithProvider(ctx context.Context, prov llm.Provider, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec { if prov == nil { return nil } diff --git a/agent/mainagent.go b/agent/mainagent.go index 6736ded..86b198b 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -21,6 +21,7 @@ type MainAgent struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) proxyAddr string // recording proxy for WebFetch (empty = direct) proxyCACert string // recording proxy's CA cert path (HTTPS verify) @@ -32,6 +33,15 @@ func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns} } +func (m *MainAgent) SetCompactionWindowResolver(fn func() int) { m.windowFn = fn } + +func (m *MainAgent) compactionWindow() int { + if m.windowFn != nil { + return m.windowFn() + } + return m.window +} + // SetProxy points the main agent's WebFetch at the recording proxy plus the CA // cert it trusts to verify HTTPS through it (empty addr = direct). func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert } @@ -70,9 +80,9 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t tsx.SetAssetStore(as, as.Companies()) } tsx.SetTaskID(taskID) - tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced) - tsx.SetResumeTask(resume) // set_goals 新增目标 → 把已完成/暂停的任务拉回 running - tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标：…」触发 + 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()...) tools, def, cleanup := AugmentTools(ctx, "mainagent", base) @@ -101,9 +111,9 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA WorkingDir: mainDir, // 本任务工作目录 /tasks/ ToolOutputDir: cmdOutDir(mainDir), - MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) - Compaction: compactionConfig(m.window), // long chats stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 + MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management) + Compaction: compactionConfig(m.compactionWindow()), // long chats stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 // 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。 Settlement: wrapupSettlement("mainagent", nil), } diff --git a/agent/planner.go b/agent/planner.go index b567d30..b2757fe 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -24,6 +24,7 @@ type Planner struct { model string tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) killWork func(intentID int64) error // engine callback to terminate a running work (nil = off) steerWork func(intentID int64, msg string) error // engine callback to steer a running work mid-run (nil = off) @@ -44,6 +45,15 @@ func NewPlanner(prov llm.Provider, model, workDir string, tx *transcript.Store, return &Planner{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, todos: map[int64]*actool.TodoStore{}} } +func (p *Planner) SetCompactionWindowResolver(fn func() int) { p.windowFn = fn } + +func (p *Planner) compactionWindow() int { + if p.windowFn != nil { + return p.windowFn() + } + return p.window +} + // SetProxy points the planner's WebFetch at the recording proxy plus the CA cert // it trusts to verify HTTPS through it (empty addr = direct). func (p *Planner) SetProxy(addr, caCert string) { p.proxyAddr, p.proxyCACert = addr, caCert } @@ -94,10 +104,11 @@ func renderPlannerTodos(items []actool.Todo) string { // TriggerEvent describes what concretely caused this planning round to fire, so // the planner looks first at the actual change instead of re-scanning the whole // overview. Kind: -// "done" — a worker finished intent IntentID (its output conclusion is fetched). -// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). -// "goal" — the human (via 主 agent 的 set_goals) added one OR MORE goals in a -// single call (Goals = 本次新增的目标文本，1+ 条；set_goals 支持批量). +// +// "done" — a worker finished intent IntentID (its output conclusion is fetched). +// "finding" — a worker reported a finding on intent IntentID (Detail = 摘要). +// "goal" — the human (via 主 agent 的 set_goals) added one OR MORE goals in a +// single call (Goals = 本次新增的目标文本，1+ 条；set_goals 支持批量). type TriggerEvent struct { Kind string IntentID int64 @@ -308,7 +319,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts ToolOutputDir: cmdOutDir(taskDir), MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management) MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余 - Compaction: compactionConfig(p.window), + Compaction: compactionConfig(p.compactionWindow()), // 跨唤醒共享的规划待办：让串行链在多轮之间保留（session 是新的，store 不是）。 Todos: p.todoFor(ts.ID()), // 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、 diff --git a/agent/provider.go b/agent/provider.go index 7a0a78e..465812b 100644 --- a/agent/provider.go +++ b/agent/provider.go @@ -8,8 +8,14 @@ package agent import ( + "bytes" "context" + "fmt" + "io" + "net/http" + "net/url" "os" + "regexp" "strings" "time" @@ -167,12 +173,16 @@ func (c Config) Provider() string { // limiter lives on the single provider instance — so planner + all workers + // main agent (which share this provider) are bounded by one shared rate limit. func (c Config) NewProvider() (llm.Provider, error) { + client, err := quotaAwareHTTPClient(c.Proxy) + if err != nil { + return nil, err + } lc := llm.Config{ - Format: c.Format, - BaseURL: c.BaseURL, - APIKey: c.APIKey, - Model: c.Model, - Proxy: c.Proxy, + Format: c.Format, + BaseURL: c.BaseURL, + APIKey: c.APIKey, + Model: c.Model, + HTTPClient: client, } // Derive the provider thinking params from the single UI selector: // "" → omit both (provider default); "off" → disable; else → enable + effort. @@ -191,6 +201,93 @@ func (c Config) NewProvider() (llm.Provider, error) { return llm.NewProvider(lc) } +// IsQuotaExhaustedMessage deliberately recognizes only explicit balance, +// billing, credit, or quota-exhaustion signals. Generic 429/rate-limit text, +// authentication failures, network errors, and server failures are excluded. +var nonFailoverHTTPStatus = regexp.MustCompile(`(?:status(?:\s+code)?|http(?:\s+status)?)\s*[=:]?\s*(?:401|403|5\d\d)\b`) +var transientQuotaLimit = regexp.MustCompile(`(?i)(?:\b(?:rpm|tpm|rpd|qps)\b|quota[_\s-]*metric|rate[_\s-]*limit|too many requests|(?:requests?|tokens?)\s+(?:per|/)\s*(?:second|minute)|(?:per|/)\s*(?:second|minute)\s+(?:requests?|tokens?)|generate[_\s-]*requests[_\s-]*per[_\s-]*(?:minute|second)|tokens?[_\s-]*per[_\s-]*(?:minute|second))`) + +func IsQuotaExhaustedMessage(message string) bool { + message = strings.ToLower(message) + // Authentication/authorization and provider-side 5xx failures never rotate, + // even when a gateway happens to echo a quota-looking phrase in the body. + if nonFailoverHTTPStatus.MatchString(message) { + return false + } + // Provider APIs frequently describe an ordinary rate limit as "quota + // exceeded", especially Google-style responses containing a quota metric. + // These limits recover with time and must stay on the current provider. + if transientQuotaLimit.MatchString(message) { + return false + } + markers := []string{ + "insufficient_quota", "quota_exceeded", "quota exceeded", "quota exhausted", + "exceeded your current quota", "billing_hard_limit_reached", + "billing hard limit", "billing_not_active", "credit balance", "insufficient credit", + "insufficient balance", "balance is too low", "payment required", "status 402", + "余额不足", "额度不足", "额度已用尽", "欠费", + } + for _, marker := range markers { + if strings.Contains(message, marker) { + return true + } + } + // gRPC RESOURCE_EXHAUSTED is overloaded for both account quota and ordinary + // request-rate limiting. Preserve it as an explicit exhaustion signal only + // when the same error does not identify a transient rate limit. + return strings.Contains(message, "resource_exhausted") && + !strings.Contains(message, "rate limit") && + !strings.Contains(message, "too many requests") +} + +// quotaAwareTransport preserves Norma's normal retry behavior except for a 429 +// whose body explicitly says the account quota/balance is exhausted. Norma's +// retry loop treats every 429 as transient; normalizing only that response to +// 402 lets a task router fail over immediately while retaining the original +// response body for provider-specific classification and audit logs. +type quotaAwareTransport struct{ base http.RoundTripper } + +func (t quotaAwareTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil || resp == nil || resp.StatusCode != http.StatusTooManyRequests { + return resp, err + } + body, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(body)) + resp.ContentLength = int64(len(body)) + if readErr != nil { + return resp, nil + } + if IsQuotaExhaustedMessage(string(body)) { + resp.StatusCode = http.StatusPaymentRequired + resp.Status = "402 Payment Required" + } + return resp, nil +} + +func quotaAwareHTTPClient(proxy string) (*http.Client, error) { + transport := http.DefaultTransport.(*http.Transport).Clone() + proxy = strings.TrimSpace(proxy) + if proxy == "" { + transport.Proxy = http.ProxyFromEnvironment + } else { + proxyURL, err := url.Parse(proxy) + if err != nil { + return nil, fmt.Errorf("llm: invalid proxy %q: %w", proxy, err) + } + switch proxyURL.Scheme { + case "http", "https", "socks5": + case "": + return nil, fmt.Errorf("llm: proxy %q missing scheme (use http://, https:// or socks5://)", proxy) + default: + return nil, fmt.Errorf("llm: unsupported proxy scheme %q (use http, https or socks5)", proxyURL.Scheme) + } + transport.Proxy = http.ProxyURL(proxyURL) + } + return &http.Client{Transport: quotaAwareTransport{base: transport}}, nil +} + // TestConnection makes a minimal real completion to verify the provider/model/ // endpoint/key actually work. Returns the round-trip latency. func TestConnection(ctx context.Context, c Config) (time.Duration, error) { diff --git a/agent/provider_quota_test.go b/agent/provider_quota_test.go new file mode 100644 index 0000000..d905ad3 --- /dev/null +++ b/agent/provider_quota_test.go @@ -0,0 +1,131 @@ +package agent + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/Autumn-27/norma/llm" +) + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestIsQuotaExhaustedMessage(t *testing.T) { + t.Parallel() + positive := []string{ + `{"error":{"code":"insufficient_quota"}}`, + `RESOURCE_EXHAUSTED`, + `You exceeded your current quota, please check your plan and billing details.`, + `billing_not_active`, + `credit balance is too low`, + `账户余额不足，请充值`, + } + for _, message := range positive { + if !IsQuotaExhaustedMessage(message) { + t.Errorf("expected quota classification for %q", message) + } + } + + negative := []string{ + `status 429: rate limit exceeded`, + `RESOURCE_EXHAUSTED: rate limit exceeded`, + `too many requests per minute`, + `quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`, + `RESOURCE_EXHAUSTED: TPM quota exceeded`, + `tokens per minute quota exceeded`, + `rate_limit_exceeded: requests per second`, + `status 401: invalid api key`, + `status 401: insufficient_quota`, + `HTTP 403: billing_hard_limit_reached`, + `status 500: internal server error`, + `status=503: insufficient_quota`, + `context length exceeded`, + } + for _, message := range negative { + if IsQuotaExhaustedMessage(message) { + t.Errorf("unexpected quota classification for %q", message) + } + } +} + +func TestQuotaAwareTransportOnlyNormalizesExplicitQuota429(t *testing.T) { + t.Parallel() + tests := []struct { + name string + status int + body string + wantStatus int + }{ + {name: "quota", status: http.StatusTooManyRequests, body: `{"code":"insufficient_quota"}`, wantStatus: http.StatusPaymentRequired}, + {name: "ordinary rate limit", status: http.StatusTooManyRequests, body: `{"message":"rate limit exceeded"}`, wantStatus: http.StatusTooManyRequests}, + {name: "server error", status: http.StatusInternalServerError, body: `insufficient_quota`, wantStatus: http.StatusInternalServerError}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + transport := quotaAwareTransport{base: roundTripperFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: tt.status, + Status: http.StatusText(tt.status), + Body: io.NopCloser(strings.NewReader(tt.body)), + Header: make(http.Header), + }, nil + })} + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.invalid", nil) + if err != nil { + t.Fatal(err) + } + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != tt.wantStatus { + t.Fatalf("status=%d, want %d", resp.StatusCode, tt.wantStatus) + } + gotBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if string(gotBody) != tt.body { + t.Fatalf("body=%q, want %q", gotBody, tt.body) + } + }) + } +} + +func TestProviderDoesNotRetryExplicitQuota429(t *testing.T) { + t.Parallel() + var requests atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + w.Header().Set("content-type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"code":"insufficient_quota","message":"You exceeded your current quota"}}`) + })) + defer upstream.Close() + + provider, err := ConfigFrom("openai", "test-model", upstream.URL, "test-key", "").NewProvider() + if err != nil { + t.Fatal(err) + } + var streamErr error + for _, err := range provider.Stream(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{llm.UserText("ping")}, + }) { + if err != nil { + streamErr = err + } + } + if streamErr == nil || !IsQuotaExhaustedMessage(streamErr.Error()) { + t.Fatalf("expected explicit quota error, got %v", streamErr) + } + if got := requests.Load(); got != 1 { + t.Fatalf("explicit quota request retried %d times, want exactly one request", got) + } +} diff --git a/agent/terminalreason.go b/agent/terminalreason.go index 90228b2..838cbd1 100644 --- a/agent/terminalreason.go +++ b/agent/terminalreason.go @@ -9,26 +9,15 @@ import ( "github.com/Autumn-27/norma/harness" ) -// This file turns a run's terminal state into something a human can act on. -// -// Before: a cancelled run produced one fixed sentence — "（运行被中断：任务暂停 / -// 终止 / 重启所致；未完成）" — a guess-list, not a diagnosis. The information to do -// better was already there (the cancel cause, the terminal reason, the turn -// count, the tool that never returned); it was just being thrown away. -// -// terminalText produces a self-contained ONE-LINE summary (the transcript's -// collapsed row, cut at the first newline by firstLine) plus a multi-line detail -// block rendered in full when the row is expanded. - -// runTrace records which tool call is in flight, so an aborted / timed-out run can -// name the step it died on — the single most useful fact when a run hangs. +// 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 // tool_use id of the most recent call - name string // its tool name - input string // its raw input JSON + id string + name string + input string at time.Time - pending bool // true = no tool_result arrived for it yet + pending bool } func (t *runTrace) start(id, name, input string) { @@ -41,73 +30,73 @@ func (t *runTrace) done(id string) { } } -// reasonHint explains each harness terminal reason in plain 中文: what it means, -// whether results were still written back, and what the engine does next. Covers -// every TerminalReason the harness can emit. var reasonHint = map[harness.TerminalReason]string{ - harness.ReasonCompleted: "模型主动结束了本轮（正常收场），但没有留下文字总结。事实/资产以工具调用为准，看上面的 tool_use 记录", - harness.ReasonMaxTurns: "撞到步数上限(MaxTurns)：SDK 已自动跑了一轮收尾(settlement)把事实/资产写回，只是缺一段文字总结。" + - "意图标记为 exhausted——表示「这个方向试过但没做完」，规划者会据此换角度，不是失败", - harness.ReasonTimeout: "撞到单次运行的软墙钟预算(MaxDuration)：SDK 在回合边界优雅收尾并写回事实/资产，只是缺文字总结。" + - "意图标记为 exhausted，不是失败", - harness.ReasonModelError: "模型/API 调用失败（网络、鉴权、限流、供应商 5xx 等）。引擎会退避后自动重跑几次；重试用尽仍失败则意图标记为 blocked", - harness.ReasonBlockingLimit: "上下文长度撞硬上限，请求在发出前就被拦下。通常是单次运行塞进了过多工具输出——考虑收窄意图粒度或压缩工具返回", - harness.ReasonPromptTooLong: "提示词过长且压缩(compaction)重试已用尽，无法继续。同样指向单次运行累积的上下文过大", - harness.ReasonImageError: "多模态内容不被当前模型支持（截图/图片块）。换支持视觉的模型，或让工具别回图", - harness.ReasonStopHookPrevented: "Stop 钩子阻止了本轮结束（守卫判定还不能收场），且随后没能继续下去。" + - "检查任务的 Guard 规则是否过严", - harness.ReasonHookStopped: "某个工具/钩子主动叫停了继续执行（守卫拦截，如越界目标、禁用命令）。看上面最后一条 tool_result 里的拦截说明", - harness.ReasonAbortedStreaming: "运行在【模型输出流式生成阶段】被取消", - harness.ReasonAbortedTools: "运行在【工具执行阶段】被取消", + 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 (summary, detail). -// summary is one self-contained line; detail is the expandable breakdown. +// 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 + } - // ---- one-line summary: cause first (that's what the reader is after), then - // just enough run shape to judge it without expanding. ---- var sum string - switch { - case aborted: + if aborted { _, short, _, ok := AbortReason(ctx) if !ok { - short = "取不到取消原因（run 的 context 上没有挂 cause）" + short = "未能取得取消原因" } - stage := "工具执行阶段" - if reason == harness.ReasonAbortedStreaming { + stage := "执行过程中" + switch reason { + case harness.ReasonAbortedStreaming: stage = "模型输出阶段" + case harness.ReasonAbortedTools: + stage = "工具执行阶段" } sum = "（运行被中断：" + short + "；停在" + stage + progressSuffix(term, tr) + "，未完成）" - case reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout: - sum = "（达运行预算上限(" + string(reason) + ")，已收尾写回事实" + progressSuffix(term, tr) + "；本次无文字总结）" - default: - hint := reasonHint[reason] - if hint == "" { - hint = "未知终态（harness 新增了终态但这里还没登记，请补 reasonHint）" - } - sum = "（无文字总结，终态 " + string(reason) + "：" + firstLine(hint, 80) + "）" + } else if reason == harness.ReasonMaxTurns || reason == harness.ReasonTimeout { + sum = "（达到运行预算上限(" + string(reason) + ")，已收尾写回事实" + progressSuffix(term, tr) + "；本次无文字总结）" + } else { + hint := terminalReasonHint(reason) + sum = "（无文字总结，终态 " + terminalReasonLabel(reason) + "：" + firstLine(hint, 80) + "）" } - // ---- expandable detail ---- var b strings.Builder b.WriteString(sum) b.WriteString("\n\n") - fmt.Fprintf(&b, "- **终态**: `%s` — %s\n", reason, reasonHint[reason]) + 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 挂上具名原因。" + - "这属于漏网的取消点，请补 agent/cancelcause.go 里的枚举\n") + 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) } @@ -118,21 +107,34 @@ func terminalText(ctx context.Context, term *harness.Terminal, tr *runTrace) (st fmt.Fprintf(&b, "- **累计 token**: 输入 %d / 输出 %d / 缓存读 %d / 缓存写 %d\n", u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheWriteTokens) } - if tr.name != "" { - if tr.pending { - fmt.Fprintf(&b, "- **中断时正在执行的工具**: `%s`(已跑 %s，**未返回结果**) — 若是硬超时/长时间无进展，问题多半就在这一步\n ```\n %s\n ```\n", - tr.name, roundDur(time.Since(tr.at)), firstLine(tr.input, 300)) - } else { - fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`(已正常返回)\n", tr.name) - } + 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 { - b.WriteString("- **工具调用**: 本次运行还没发出任何工具调用就结束了\n") + fmt.Fprintf(&b, "- **中断前最后一个工具**: `%s`（已正常返回）\n", tr.name) } return sum, b.String() } -// progressSuffix adds "已跑 N 轮 / 12m30s" to a summary line — how far the run got -// before it died, which decides whether an abort cost anything worth re-running. +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 { @@ -144,10 +146,9 @@ func progressSuffix(term *harness.Terminal, tr *runTrace) string { if len(parts) == 0 { return "" } - return "，已跑 " + strings.Join(parts, " / ") + return "，已运行 " + strings.Join(parts, " / ") } -// roundDur formats a duration at a readable granularity (no 3.400213s noise). func roundDur(d time.Duration) string { switch { case d < time.Minute: diff --git a/agent/terminalreason_test.go b/agent/terminalreason_test.go index a3ecd20..ff53fe4 100644 --- a/agent/terminalreason_test.go +++ b/agent/terminalreason_test.go @@ -12,8 +12,6 @@ import ( "github.com/Autumn-27/norma/llm" ) -// every TerminalReason the harness can emit must have an explanation — otherwise -// a run silently falls back to "未知终态", which is the bug this file exists to fix. func TestReasonHintCoversEveryTerminalReason(t *testing.T) { all := []harness.TerminalReason{ harness.ReasonCompleted, harness.ReasonBlockingLimit, harness.ReasonImageError, @@ -21,152 +19,151 @@ func TestReasonHintCoversEveryTerminalReason(t *testing.T) { harness.ReasonPromptTooLong, harness.ReasonStopHookPrevented, harness.ReasonHookStopped, harness.ReasonMaxTurns, harness.ReasonTimeout, } - for _, r := range all { - if strings.TrimSpace(reasonHint[r]) == "" { - t.Errorf("terminal reason %q 没有中文说明", r) + for _, reason := range all { + if strings.TrimSpace(reasonHint[reason]) == "" { + t.Errorf("terminal reason %q has no explanation", reason) } } } -// the cause attached at cancel time must survive the ctx chain the engine builds: -// task exec ctx -> per-work child -> value ctx -> the run's hard-timeout backstop. func TestAbortCausePropagatesThroughRunContextChain(t *testing.T) { type key struct{} - exec, cancelExec := context.WithCancelCause(context.Background()) - work, cancelWork := context.WithCancelCause(exec) + execCtx, cancelExec := context.WithCancelCause(context.Background()) + workCtx, cancelWork := context.WithCancelCause(execCtx) defer cancelWork(nil) - valued := context.WithValue(work, key{}, "task-1") - run, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) + valued := context.WithValue(workCtx, key{}, "task-1") + runCtx, cancelRun := context.WithTimeoutCause(valued, time.Hour, AbortRunHardTimeout) defer cancelRun() cancelExec(AbortPausedByUser) - - code, _, text, ok := AbortReason(run) + code, _, text, ok := AbortReason(runCtx) if !ok || code != "paused_by_user" { - t.Fatalf("code=%q ok=%v，期望 paused_by_user", code, ok) + t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok) } - if !strings.Contains(text, "退回 frontier") { - t.Errorf("原因文案没传下来: %q", text) + 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.Error("未取消的 ctx 不应报出中断原因") + t.Fatal("live context must not report an abort reason") } - // a cancel site that forgot to attach a cause must be called out, not guessed at plain, cancel := context.WithCancel(context.Background()) cancel() - code, _, _, ok := AbortReason(plain) - if !ok || code != "canceled_no_cause" { - t.Errorf("code=%q ok=%v，期望 canceled_no_cause", code, ok) + if code, _, _, ok := AbortReason(plain); !ok || code != "canceled_no_cause" { + t.Fatalf("code=%q ok=%v, want canceled_no_cause", code, ok) } - - timed, cancelT := context.WithTimeout(context.Background(), time.Nanosecond) - defer cancelT() + timed, cancelTimed := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancelTimed() <-timed.Done() if code, _, _, _ := AbortReason(timed); code != "deadline_exceeded" { - t.Errorf("code=%q，期望 deadline_exceeded", code) + t.Fatalf("code=%q, want deadline_exceeded", code) } } -// an aborted run's one-line summary must name the actual cause, and the detail -// must pin down the tool that never came back. func TestTerminalTextAbortedNamesCauseAndHangingTool(t *testing.T) { ctx, cancel := context.WithCancelCause(context.Background()) cancel(AbortKilledByPlanner) - - tr := &runTrace{startedAt: time.Now().Add(-90 * time.Second)} - tr.start("tu_1", "Bash", `{"command":"nmap -p- 10.0.0.1"}`) - + 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}, } - sum, detail := terminalText(ctx, term, tr) - - if !strings.Contains(sum, "kill_work") { - t.Errorf("摘要没点出取消方: %q", sum) - } - if strings.Contains(sum, "\n") { - t.Errorf("摘要必须是单行（前端按首行截断）: %q", sum) + 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("详情缺少 %q:\n%s", want, detail) + t.Errorf("detail missing %q:\n%s", want, detail) } } } -// a tool that DID return must not be reported as the hang site. +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) - tr := &runTrace{startedAt: time.Now()} - tr.start("tu_1", "Read", `{"path":"/etc/hosts"}`) - tr.done("tu_1") - - _, detail := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonAbortedStreaming}, tr) - if strings.Contains(detail, "未返回结果") { - t.Errorf("已返回的工具被误报为卡住:\n%s", detail) - } - if !strings.Contains(detail, "已正常返回") { - t.Errorf("详情没记录最后一个工具:\n%s", detail) + 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) } } -// non-abort terminals keep their own explanation instead of the interruption text. func TestTerminalTextNonAbortReasons(t *testing.T) { - ctx := context.Background() - tr := &runTrace{startedAt: time.Now()} - - sum, _ := terminalText(ctx, &harness.Terminal{Reason: harness.ReasonMaxTurns}, tr) - if !strings.Contains(sum, "达运行预算上限") || strings.Contains(sum, "中断") { - t.Errorf("max_turns 摘要不对: %q", sum) + 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) } - - sum, detail := terminalText(ctx, &harness.Terminal{ - Reason: harness.ReasonModelError, Err: errors.New("429 rate limited")}, tr) - if !strings.Contains(sum, "model_error") { - t.Errorf("model_error 摘要不对: %q", sum) - } - if !strings.Contains(detail, "429 rate limited") { - t.Errorf("详情没带上底层错误:\n%s", detail) + 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) } } -// every enumerated cause must carry both a compact headline (for the collapsed -// row) and a full explanation (for the expanded detail). func TestAbortCausesAreWellFormed(t *testing.T) { all := []*AbortCause{ AbortPausedByUser, AbortPausedByOrchestrator, AbortTaskDeleted, AbortPausedOnReload, - AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkFinished, - AbortPausedRaceGuard, AbortChatStoppedByUser, AbortChatTurnFinished, + AbortGoalMet, AbortSettleDrainTimeout, AbortKilledByPlanner, AbortWorkPausedByUser, + AbortWorkCancelledByUser, AbortWorkFinished, AbortPausedRaceGuard, + AbortChatStoppedByUser, AbortChatPausedWithTask, AbortChatTurnFinished, AbortShutdown, AbortRunHardTimeout, } seen := map[string]bool{} - for _, c := range all { + for _, abort := range all { switch { - case c.Code == "" || seen[c.Code]: - t.Errorf("code 缺失或重复: %q", c.Code) - case c.Short == "" || utf8.RuneCountInString(c.Short) > 40: - t.Errorf("%s: Short 必须存在且 ≤40 字（进折叠摘要行）: %q", c.Code, c.Short) - case utf8.RuneCountInString(c.Text) <= utf8.RuneCountInString(c.Short): - t.Errorf("%s: Text 应比 Short 更完整", c.Code) + 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[c.Code] = true + seen[abort.Code] = true } } -// summaries are mostly 中文; capping by bytes used to slice a rune in half. func TestFirstLineCapsByRuneNotByte(t *testing.T) { - got := firstLine(strings.Repeat("恢复", 10), 5) - if got != "恢复恢复恢…" { - t.Errorf("firstLine 按字节截断了: %q", got) + if got := firstLine(strings.Repeat("恢复", 10), 5); got != "恢复恢复恢…" { + t.Fatalf("firstLine split a Unicode character: %q", got) } if got := firstLine("头\n尾", 100); got != "头" { - t.Errorf("多行没截到首行: %q", got) + t.Fatalf("firstLine did not stop at newline: %q", got) } } diff --git a/agent/tools.go b/agent/tools.go index 93b71b7..d0958a2 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -23,6 +23,10 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[ var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "summary": p["summary"], "state": n.State} + if n.Inherited { + m["source_task_id"] = n.SourceTaskID + m["inherited"] = true + } // asset_ids is the structured "which assets this direction covers" signal for // dedup; fall back to legacy payload keys (target_ids plural, then target_id // single) so intents stored before the rename still surface their anchors. @@ -265,132 +269,438 @@ func (t *ToolSet) graphOverview() actool.CoreTool { // the model needn't spend a turn calling the tool — every plan round starts with // an empty context and always needs this first). func (t *ToolSet) graphOverviewData() map[string]any { - { - out := map[string]any{} - // goals summary folded in so the planner needn't call list_goals each round. - goals, _ := t.ts.ListByKind(db.KindGoal, 100) - gsum := make([]map[string]any, 0, len(goals)) - for _, g := range goals { - var p map[string]any - _ = json.Unmarshal(g.Payload, &p) - gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + out := map[string]any{} + // goals summary folded in so the planner needn't call list_goals each round. + goals, _ := t.ts.ListByKind(db.KindGoal, 100) + gsum := make([]map[string]any, 0, len(goals)) + for _, g := range goals { + var p map[string]any + _ = json.Unmarshal(g.Payload, &p) + gsum = append(gsum, map[string]any{"id": g.ID, "state": g.State, "text": p["text"]}) + } + out["goals"] = gsum + // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示；folded in so the + // planner reads them every round when generating intents (否则只写不读). + hints, _ := t.ts.ListByKind(db.KindHint, 50) + hsum := make([]map[string]any, 0, len(hints)) + for _, h := range hints { + var p map[string]any + _ = json.Unmarshal(h.Payload, &p) + hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + out["hints"] = hsum + // lineage from the exploration edges: an intent's parents (what it + // derived_from — possibly several facts combined) and its yields (the + // facts/findings it produced). factFrom maps a fact → the intent that + // produced it. This is the relationship layer the flat lists lacked. + edges, _ := t.ts.Edges(5000) + parentsOf := map[int64][]int64{} + yieldsOf := map[int64][]int64{} + factFrom := map[int64]int64{} + for _, e := range edges { + switch e.Rel { + case db.RelDerivedFrom, db.RelSpawns: // upstream: derived_from (fact/finding/intent→intent) or spawns (origin fact→goal, legacy begin→intent) + parentsOf[e.To] = append(parentsOf[e.To], e.From) + case db.RelYields: // intent --yields--> fact/finding + yieldsOf[e.From] = append(yieldsOf[e.From], e.To) + factFrom[e.To] = e.From } - out["goals"] = gsum - // hints: 人类/主 agent 通过 add_hint 挂上图的战略提示；folded in so the - // planner reads them every round when generating intents (否则只写不读). - hints, _ := t.ts.ListByKind(db.KindHint, 50) - hsum := make([]map[string]any, 0, len(hints)) - for _, h := range hints { - var p map[string]any - _ = json.Unmarshal(h.Payload, &p) - hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}) + } + fr, _ := t.ts.Frontier(100) + out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) + all, _ := t.ts.ListByKind(db.KindIntent, 300) + var running, recentDone []*db.Node + for _, n := range all { + switch n.State { + case "running": + running = append(running, n) + case "done", "blocked", "exhausted": + if len(recentDone) < 15 { + recentDone = append(recentDone, n) + } } - out["hints"] = hsum - // lineage from the exploration edges: an intent's parents (what it - // derived_from — possibly several facts combined) and its yields (the - // facts/findings it produced). factFrom maps a fact → the intent that - // produced it. This is the relationship layer the flat lists lacked. - edges, _ := t.ts.Edges(5000) + } + out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) + out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) + out["frontier_open"] = len(fr) + // findings (confirmed vulns) and facts (worker exploration results) are + // now distinct node kinds. recent_facts surfaces fact summaries (esp. + // negative results) so the planner sees them in one call; full content + // via node_detail(id). + vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000) + factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first + out["findings"] = len(vulnNodes) // 确认漏洞数（目标判定看它） + out["facts"] = len(factNodes) // 探索事实/结论数（含否定结论） + recentFacts := make([]map[string]any, 0, 20) + for _, n := range factNodes { + if len(recentFacts) >= 20 { + break + } + m := compactNode(n) + if from := factFrom[n.ID]; from > 0 { + m["from_intent"] = from // 本事实由哪个意图产生 + } + // confidence 带进概览：让规划者一眼看出哪条结论只是 inferred（尤其否定结论 + // 别当铁案）；evidence 较长，留给 node_detail(id)。 + var fp map[string]any + if json.Unmarshal(n.Payload, &fp) == nil { + if c, ok := fp["confidence"].(string); ok && c != "" { + m["confidence"] = c + } + } + recentFacts = append(recentFacts, m) + } + out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?}，详情/证据用 node_detail(id) + // the original task (root) so the planner always has it, not just the + // decomposed goals. + if description, goal, err := t.ts.Root(); err == nil { + out["task"] = map[string]any{"description": description, "goal": goal} + } + // Direct source tasks are a live, read-only blackboard view. Keep their + // summaries in a separate field so their intents never enter this task's + // frontier or get mistaken for locally claimable work. + out["related_tasks"] = t.relatedTaskOverviews() + // coverage：粗略的资产测试覆盖度参考——范围(task_scope)内的资产里，被 fact 碰过的 + // 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。 + if t.as != nil && t.ts != nil && t.taskID > 0 { + if cov, err := t.as.TaskCoverageWithSources(t.taskID); err == nil { + m := map[string]any{ + "denominator": cov.Denominator, + "tested": cov.Tested, + "by_type": cov.ByType, + "note": "coverage资产测试覆盖度（包括接口等各种相关资产），粗略估计、仅供参考：包含当前任务与直接关联任务的 scope、事实锚点；关联 scope 只读。容器型资产/大量枚举会让它偏低，勿据此认为已测完；可用 add_task_scope 增补本任务范围、list_untested_assets 看未测资产【通常不调用list_untested_assets，按照任务推进即可】；", + } + if cov.Denominator == 0 { + m["pct"] = nil + m["status"] = "范围未锚定" + } else { + m["pct"] = cov.Pct + } + // scope：当前测试范围的根资产（task_scope 原始行），让 agent 知道这个任务到底 + // 圈定了哪些目标（不是全部测试资产，而是范围边界本身）。 + if rows, err := t.as.ListTaskScopeWithSources(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, "task_id": r.TaskID} + if r.TaskID != t.taskID { + inheritedMap(e, r.TaskID) + } + 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 + } + if hosts, err := t.as.HostsByTaskWithSources(t.taskID); err == nil { + const hostContextLimit = 500 + visible := hosts + if len(visible) > hostContextLimit { + visible = visible[:hostContextLimit] + } + m["hosts"] = visible + m["host_count"] = len(hosts) + if len(visible) < len(hosts) { + m["hosts_truncated"] = true + } + } + out["coverage"] = m + } + } + return out +} + +func inheritedMap(m map[string]any, sourceTaskID int64) map[string]any { + m["source_task_id"] = sourceTaskID + m["inherited"] = true + return m +} + +const ( + relatedOverviewTotalTextRunes = 48_000 + relatedOverviewMaxTextPerSource = 8_000 + relatedOverviewMaxGoalsPerSource = 8 + relatedOverviewMaxHintsPerSource = 6 + relatedOverviewMaxFactsPerSource = 12 + relatedOverviewMaxFindingsPerTask = 6 + relatedOverviewMaxIntentsPerTask = 8 + relatedOverviewMaxScopePerSource = 12 +) + +// overviewTextBudget bounds inherited prompt text while preserving a fair slice +// for every direct source. Full evidence remains available through the on-demand +// read tools, so truncation here does not discard persisted blackboard data. +type overviewTextBudget struct { + remaining int + truncated bool +} + +func relatedOverviewBudgetForSources(sourceCount int) int { + if sourceCount <= 0 { + return 0 + } + if sourceCount > db.MaxTaskSourceCount { + sourceCount = db.MaxTaskSourceCount + } + perSource := relatedOverviewTotalTextRunes / sourceCount + if perSource > relatedOverviewMaxTextPerSource { + perSource = relatedOverviewMaxTextPerSource + } + return perSource +} + +func (b *overviewTextBudget) take(value any, fieldLimit int) string { + var text string + switch value := value.(type) { + case string: + text = strings.TrimSpace(value) + case nil: + return "" + default: + text = strings.TrimSpace(fmt.Sprint(value)) + } + if text == "" { + return "" + } + if b.remaining <= 0 || fieldLimit <= 0 { + b.truncated = true + return "" + } + runes := []rune(text) + limit := fieldLimit + if limit > b.remaining { + limit = b.remaining + } + if len(runes) > limit { + b.truncated = true + if limit == 1 { + text = "…" + } else { + text = string(runes[:limit-1]) + "…" + } + runes = []rune(text) + } + b.remaining -= len(runes) + return text +} + +func recentTerminalIntents(store *db.ExplorationStore, limit int) []*db.Node { + if limit <= 0 { + return []*db.Node{} + } + const batch = 300 + cursor := int64(0) + out := make([]*db.Node, 0, limit) + for len(out) < limit { + page, more, err := store.ListByKindPage(db.KindIntent, cursor, batch) + if err != nil || len(page) == 0 { + break + } + for _, intent := range page { + switch intent.State { + case "done", "blocked", "exhausted", "stopped": + out = append(out, intent) + } + if len(out) >= limit { + break + } + } + if !more { + break + } + cursor = page[len(page)-1].ID + } + return out +} + +// relatedTaskOverviews distills persistent blackboard state from direct source +// tasks. It intentionally reads each source's local store methods, never its own +// related sources, so inheritance is one level only. +func (t *ToolSet) relatedTaskOverviews() []map[string]any { + sources, err := t.ts.DirectSourceStores() + if err != nil { + return []map[string]any{} + } + if len(sources) > db.MaxTaskSourceCount { + sources = sources[:db.MaxTaskSourceCount] + } + perSourceTextBudget := relatedOverviewBudgetForSources(len(sources)) + out := make([]map[string]any, 0, len(sources)) + for _, source := range sources { + ts := source.Store + budget := overviewTextBudget{remaining: perSourceTextBudget} + item := map[string]any{ + "source_task_id": source.Task.TaskID, + "inherited": true, + "task": map[string]any{ + "description": budget.take(source.Task.Description, 800), + "goal": budget.take(source.Task.Goal, 800), + "status": source.Task.Status, + }, + } + stats, statsErr := ts.Stats() + + edges, _ := 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 + for _, edge := range edges { + switch edge.Rel { + case db.RelDerivedFrom, db.RelSpawns: + parentsOf[edge.To] = append(parentsOf[edge.To], edge.From) + case db.RelYields: + yieldsOf[edge.From] = append(yieldsOf[edge.From], edge.To) + factFrom[edge.To] = edge.From } } - fr, _ := t.ts.Frontier(100) - out["open_intents"] = compactIntents(fr, parentsOf, yieldsOf) - all, _ := t.ts.ListByKind(db.KindIntent, 300) - var running, recentDone []*db.Node - for _, n := range all { - switch n.State { - case "running": - running = append(running, n) - case "done", "blocked", "exhausted": - if len(recentDone) < 15 { - recentDone = append(recentDone, n) - } + + goals, _ := ts.ListByKind(db.KindGoal, relatedOverviewMaxGoalsPerSource) + goalSummary := make([]map[string]any, 0, len(goals)) + for _, goal := range goals { + var payload map[string]any + _ = json.Unmarshal(goal.Payload, &payload) + goalSummary = append(goalSummary, inheritedMap(map[string]any{ + "id": goal.ID, "state": goal.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["goals"] = goalSummary + + hints, _ := ts.ListByKind(db.KindHint, relatedOverviewMaxHintsPerSource) + hintSummary := make([]map[string]any, 0, len(hints)) + for _, hint := range hints { + var payload map[string]any + _ = json.Unmarshal(hint.Payload, &payload) + hintSummary = append(hintSummary, inheritedMap(map[string]any{ + "id": hint.ID, "state": hint.State, "text": budget.take(payload["text"], 400), + }, source.Task.TaskID)) + } + item["hints"] = hintSummary + + facts, _ := ts.ListByKind(db.KindFact, relatedOverviewMaxFactsPerSource) + findings, _ := ts.ListByKind(db.KindFinding, relatedOverviewMaxFindingsPerTask) + intentNodes, _ := ts.ListByKind(db.KindIntent, 300) + terminalIntent := make(map[int64]bool, len(intentNodes)) + for _, intent := range intentNodes { + terminalIntent[intent.ID] = inheritedIntentSummaryState(intent.State) + } + item["facts"] = len(facts) + item["findings"] = len(findings) + if statsErr == nil { + item["facts"] = stats[db.KindFact] + item["findings"] = stats[db.KindFinding] + if stats[db.KindGoal] > len(goals) || stats[db.KindHint] > len(hints) || + stats[db.KindFact] > len(facts) || stats[db.KindFinding] > len(findings) { + budget.truncated = true } } - out["running_intents"] = compactIntents(running, parentsOf, yieldsOf) - out["recent_done_intents"] = compactIntents(recentDone, parentsOf, yieldsOf) - out["frontier_open"] = len(fr) - // findings (confirmed vulns) and facts (worker exploration results) are - // now distinct node kinds. recent_facts surfaces fact summaries (esp. - // negative results) so the planner sees them in one call; full content - // via node_detail(id). - vulnNodes, _ := t.ts.ListByKind(db.KindFinding, 1000) - factNodes, _ := t.ts.ListByKind(db.KindFact, 1000) // newest first - out["findings"] = len(vulnNodes) // 确认漏洞数（目标判定看它） - out["facts"] = len(factNodes) // 探索事实/结论数（含否定结论） - recentFacts := make([]map[string]any, 0, 20) - for _, n := range factNodes { - if len(recentFacts) >= 20 { - break + recentFindings := make([]map[string]any, 0, len(findings)) + for _, finding := range findings { + entry := inheritedMap(compactFinding(finding), source.Task.TaskID) + entry["summary"] = budget.take(entry["summary"], 400) + recentFindings = append(recentFindings, entry) + } + item["recent_findings"] = recentFindings + recentFacts := make([]map[string]any, 0, len(facts)) + for _, fact := range facts { + m := inheritedMap(compactNode(fact), source.Task.TaskID) + m["summary"] = budget.take(m["summary"], 400) + if from := factFrom[fact.ID]; from > 0 && terminalIntent[from] { + m["from_intent"] = from } - m := compactNode(n) - if from := factFrom[n.ID]; from > 0 { - m["from_intent"] = from // 本事实由哪个意图产生 - } - // confidence 带进概览：让规划者一眼看出哪条结论只是 inferred（尤其否定结论 - // 别当铁案）；evidence 较长，留给 node_detail(id)。 - var fp map[string]any - if json.Unmarshal(n.Payload, &fp) == nil { - if c, ok := fp["confidence"].(string); ok && c != "" { - m["confidence"] = c + var payload map[string]any + if json.Unmarshal(fact.Payload, &payload) == nil { + if confidence, ok := payload["confidence"].(string); ok && confidence != "" { + m["confidence"] = confidence } } recentFacts = append(recentFacts, m) } - out["recent_facts"] = recentFacts // 最近事实的 {id, summary, from_intent, confidence?}，详情/证据用 node_detail(id) - // the original task (root) so the planner always has it, not just the - // decomposed goals. - if description, goal, err := t.ts.Root(); err == nil { - out["task"] = map[string]any{"description": description, "goal": goal} + item["recent_facts"] = recentFacts + + recentDone := recentTerminalIntents(ts, relatedOverviewMaxIntentsPerTask) + for _, intent := range recentDone { + intent.Inherited = true + intent.SourceTaskID = source.Task.TaskID } - // coverage：粗略的资产测试覆盖度参考——范围(task_scope)内的资产里，被 fact 碰过的 - // 占比 + by_type(按类型的 总数/已测)。要看未测的具体资产由 agent 按需调 list_untested_assets 自行判断。仅任务上下文有。 - if t.as != nil && t.ts != nil && t.taskID > 0 { - if cov, err := t.as.TaskCoverage(t.taskID, t.ts.ID()); err == nil { - m := map[string]any{ - "denominator": cov.Denominator, - "tested": cov.Tested, - "by_type": cov.ByType, - "note": "coverage资产测试覆盖度（包括接口等各种相关资产），粗略估计、仅供参考：容器型资产/大量枚举会让它偏低，勿据此认为已测完；scope 是当前任务测试范围的根资产（root_domain/subdomain/ip/cidr/company 原始行），即覆盖度分母的边界、本任务授权/圈定的目标本身；可用 add_task_scope 增补范围、list_untested_assets 看未测的具体资产【通常不调用list_untested_assets，按照任务推进即可】；", + intentResults := compactIntents(recentDone, parentsOf, yieldsOf) + for i, intent := range recentDone { + intentResults[i]["summary"] = budget.take(intentResults[i]["summary"], 400) + acts, _, err := ts.ActivityPageForTerminalIntent(intent.ID, 0, 20) + if err != nil { + continue + } + var resultSummary, textFallback string + for _, activity := range acts { + switch activity.Kind { + case "result": + resultSummary = activity.Summary + case "text": + textFallback = activity.Summary } - if cov.ScopeRows == 0 { - m["pct"] = nil - m["status"] = "范围未锚定" - } else { - m["pct"] = cov.Pct - } - // scope：当前测试范围的根资产（task_scope 原始行），让 agent 知道这个任务到底 - // 圈定了哪些目标（不是全部测试资产，而是范围边界本身）。 - if rows, err := t.as.ListTaskScope(t.taskID); err == nil && len(rows) > 0 { - scope := make([]map[string]any, 0, len(rows)) - for _, r := range rows { - e := map[string]any{"kind": r.Kind, "source": r.Source} - switch { - case r.Domain != "": - e["value"] = r.Domain - case r.Net != "": - e["value"] = r.Net - case r.CompanyID != nil: - e["company_id"] = *r.CompanyID - } - scope = append(scope, e) - } - m["scope"] = scope - } - out["coverage"] = m + } + if resultSummary == "" { + resultSummary = textFallback + } + if resultSummary != "" { + intentResults[i]["result_summary"] = budget.take(resultSummary, 800) } } - return out + item["recent_intent_results"] = intentResults + if statsErr == nil { + item["node_stats"] = stats + } + + if t.as != nil { + if scopeRows, err := t.as.ListTaskScope(source.Task.TaskID); err == nil && len(scopeRows) > 0 { + scopeCount := len(scopeRows) + if len(scopeRows) > relatedOverviewMaxScopePerSource { + scopeRows = scopeRows[:relatedOverviewMaxScopePerSource] + budget.truncated = true + } + scope := make([]map[string]any, 0, len(scopeRows)) + for _, row := range scopeRows { + entry := map[string]any{"kind": row.Kind, "source": budget.take(row.Source, 300)} + switch { + case row.Domain != "": + entry["value"] = budget.take(row.Domain, 400) + case row.Net != "": + entry["value"] = budget.take(row.Net, 400) + case row.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 } } @@ -399,7 +709,11 @@ func (t *ToolSet) graphOverviewData() map[string]any { func compactNode(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) - return map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } + return m } // compactFinding is compactNode plus the vuln-specific vulnclass/severity. @@ -407,6 +721,9 @@ func compactFinding(n *db.Node) map[string]any { var p map[string]any _ = json.Unmarshal(n.Payload, &p) m := map[string]any{"id": n.ID, "state": n.State, "summary": p["summary"]} + if n.Inherited { + inheritedMap(m, n.SourceTaskID) + } if vc, ok := p["vulnclass"]; ok && vc != nil && vc != "" { m["vulnclass"] = vc } @@ -417,16 +734,23 @@ func compactFinding(n *db.Node) map[string]any { } func (t *ToolSet) listFindings() actool.CoreTool { - return readTool("list_findings", "列【确认漏洞】(紧凑：id+task_id+intent_id+vulnclass+severity+摘要+状态)。task_id=该漏洞所属任务, intent_id=产生它的意图。这里只有漏洞,不含普通探索事实(那是 list_facts)。详情用 node_detail(id)。", + return readTool("list_findings", "列本任务及直接关联任务的【确认漏洞】(紧凑：id+task_id+intent_id+vulnclass+severity+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。这里只含漏洞；普通探索事实用 list_facts，详情用 node_detail(id)。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFinding, 500) - intentOf, _ := t.ts.FindingIntents() // finding id -> 产生它的 intent id - taskID := t.ts.ID() // exploration id = 任务 id（本 store 内所有 finding 同属） + f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500) + intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id + taskID := t.taskID + if taskID <= 0 { + taskID, _ = t.ts.TaskID() + } out := make([]map[string]any, 0, len(f)) for _, n := range f { m := compactFinding(n) - m["task_id"] = taskID + if n.Inherited { + m["task_id"] = n.SourceTaskID + } else { + m["task_id"] = taskID + } if iid, ok := intentOf[n.ID]; ok { m["intent_id"] = iid } @@ -437,10 +761,10 @@ func (t *ToolSet) listFindings() actool.CoreTool { } func (t *ToolSet) listFacts() actool.CoreTool { - return readTool("list_facts", "列【探索事实/结论】(紧凑：id+摘要+状态),即 worker 按意图探出的结论(含'端口关闭/不可注入'等否定结论)。详情用 node_detail(id)。漏洞看 list_findings。", + return readTool("list_facts", "列本任务及直接关联任务的【探索事实/结论】(紧凑：id+摘要+状态)。关联任务条目带 source_task_id/inherited=true 且只读。详情用 node_detail(id)，漏洞看 list_findings。", obj(map[string]any{}), func(context.Context, json.RawMessage) (actool.Result, error) { - f, _ := t.ts.ListByKind(db.KindFact, 500) + f, _ := t.ts.ListByKindWithSources(db.KindFact, 500) out := make([]map[string]any, 0, len(f)) for _, n := range f { out = append(out, compactNode(n)) @@ -450,7 +774,7 @@ func (t *ToolSet) listFacts() actool.CoreTool { } func (t *ToolSet) nodeDetail() actool.CoreTool { - return readTool("node_detail", "按 id 取某个【探索图节点】的完整内容(事实/发现/意图/目标：摘要 + 详情/证据/PoC)。仅限探索图节点 id(来自 list_facts / list_findings / graph_overview 的 recent_facts 返回的 id)。查【资产】请用 list_assets / asset_neighbors——资产 id 不是探索节点 id,不要传进来。", + return readTool("node_detail", "按 id 取本任务或直接关联任务的【探索图节点】完整内容。继承节点带 source_task_id/inherited=true 且只读。仅限 list_facts/list_findings/graph_overview 返回的探索节点 id；资产请用 list_assets/asset_neighbors。", obj(map[string]any{"id": idp("探索图节点 id(非资产 id)")}, "id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -461,7 +785,7 @@ func (t *ToolSet) nodeDetail() actool.CoreTool { if id <= 0 { return actool.Errorf("id 必填"), nil } - n, err := t.ts.GetNode(id) + n, err := t.ts.GetNodeWithSources(id) if err != nil { return actool.Errorf(err.Error()), nil } @@ -493,9 +817,9 @@ func (t *ToolSet) addOneIntent(it intentItem) (int64, error) { // 先校验锚点（建节点前，避免坏锚点留下孤儿意图）。 parents := pidList(it.ParentIDs) for _, pidv := range parents { - n, err := t.ts.GetNode(pidv) + n, err := t.ts.GetNodeWithSources(pidv) if err != nil || n == nil { - return 0, fmt.Errorf("parent_id %d 不存在：parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id；顶层全新方向请留空 parent_ids", pidv) + return 0, fmt.Errorf("parent_id %d 不存在于本任务或直接关联任务：parent_ids 必须是已存在的【事实(fact)/发现(finding)】节点 id；顶层全新方向请留空 parent_ids", pidv) } if n.Kind != db.KindFact && n.Kind != db.KindFinding { return 0, fmt.Errorf("parent_id %d 是 %q 节点，不能作为意图锚点：意图只能锚在已确认的【事实(fact)/发现(finding)】上，不能挂在意图/目标/提示上；顶层全新方向请留空 parent_ids", pidv, n.Kind) @@ -604,6 +928,14 @@ func (t *ToolSet) proveGoal() actool.CoreTool { if goal == 0 || ev == 0 { return actool.Errorf("goal_id 和 evidence_id 必填"), nil } + goalNode, err := t.ts.GetNode(goal) + if err != nil || goalNode == nil || goalNode.Kind != db.KindGoal { + return actool.Errorf("goal_id 必须是本任务的目标节点（关联任务目标只读）"), nil + } + evidenceNode, err := t.ts.GetNodeWithSources(ev) + if err != nil || evidenceNode == nil || (evidenceNode.Kind != db.KindFact && evidenceNode.Kind != db.KindFinding) { + return actool.Errorf("evidence_id 必须是本任务或直接关联任务的事实/漏洞节点"), nil + } _ = t.ts.Link(ev, db.RelProves, goal) _ = t.ts.SetNodeState(goal, "met") // 每标记一个目标 met，就检查本任务是否【所有目标】都已 met；若是，自动判定 @@ -668,12 +1000,19 @@ func (t *ToolSet) addFinding() actool.CoreTool { } var id int64 if t.ts != nil { + intent := pid(a.IntentID) + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } + } var err error id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors) if err != nil { return actool.Errorf(err.Error()), nil } - if intent := pid(a.IntentID); intent > 0 { + if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding } _, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors) @@ -728,15 +1067,21 @@ func (t *ToolSet) recordOneFact(it factItem, defaultIntent int64) (int64, error) if c := strings.TrimSpace(it.Confidence); c != "" { payload["confidence"] = c } + intent := pid(it.IntentID) + if intent <= 0 { + intent = defaultIntent + } + if intent > 0 { + node, err := t.ts.GetNode(intent) + if err != nil || node == nil || node.Kind != db.KindIntent { + return 0, fmt.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）") + } + } // a fact is its OWN node kind (distinct from a vuln finding). id, err := t.ts.AddNode(db.KindFact, payload, 5, "confirmed", t.worker, pidList(it.AssetIDs)) if err != nil { return 0, err } - intent := pid(it.IntentID) - if intent <= 0 { - intent = defaultIntent - } if intent > 0 { _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> fact } @@ -984,6 +1329,10 @@ func (t *ToolSet) killWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } if err := t.killWork(id); err != nil { return actool.Errorf(err.Error()), nil } @@ -1014,6 +1363,10 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + node, err := t.ts.GetNode(id) + if err != nil || node == nil || node.Kind != db.KindIntent { + return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil + } if strings.TrimSpace(a.Message) == "" { return actool.Errorf("message 必填"), nil } @@ -1026,7 +1379,7 @@ func (t *ToolSet) steerWorkTool() actool.CoreTool { // getWorkerOutput returns a work's final (or截至中止时的) conclusion text by intent id. func (t *ToolSet) getWorkerOutput() actool.CoreTool { - return readTool("get_worker_output", "取某条意图(work)的最终输出结论。正常结束返回其总结；被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", + return readTool("get_worker_output", "取本任务或直接关联任务某条意图(work)的最终输出结论。关联任务结果带 source_task_id/inherited=true 且只读。正常结束返回其总结；被终止(stopped)/异常的 work 返回其截至中止时的最后输出(terminated=true)。", obj(map[string]any{"intent_id": idp("意图 id（= work 句柄）")}, "intent_id"), func(_ context.Context, in json.RawMessage) (actool.Result, error) { var a struct { @@ -1037,7 +1390,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } - acts, _, err := t.ts.ActivityList(&id, 0, 1000) + intentNode, err := t.ts.GetNodeWithSources(id) + if err != nil { + return actool.Errorf(err.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } + acts, _, err := t.ts.ActivityListWithSources(id, 0, 1000) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1056,16 +1416,25 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { pick, terminated = fallback, true } if pick == nil { + if intentNode.Inherited { + return jsonResult(inheritedMap(map[string]any{ + "intent_id": id, "final_text": "（该 work 尚无任何输出）", "terminated": true, + }, intentNode.SourceTaskID)) + } return actool.Text("（该 work 尚无任何输出）"), nil } - detail, _ := t.ts.ActivityDetail(pick.ID) + detail, _ := t.ts.ActivityDetailWithSources(pick.ID) if detail == "" { detail = pick.Summary } - return jsonResult(map[string]any{ + result := map[string]any{ "intent_id": id, "worker_name": pick.Worker, "final_text": detail, "summary": pick.Summary, "is_error": pick.IsError, "terminated": terminated, - }) + } + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1075,10 +1444,14 @@ func (t *ToolSet) getWorkerOutput() actool.CoreTool { func traceSteps(acts []db.Activity) []map[string]any { steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } return steps } @@ -1092,7 +1465,7 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { "① 只传 intent_id → 返回该 work 每一步的摘要流（summary≤100字，含 step_id；只是动作轮廓，不含完整输出）；\n"+ "② intent_id + q → 只返回命中关键字的步骤摘要（在摘要和完整输出里都搜；仍只给 summary，要看内容用③）；\n"+ "③ intent_id + step_ids（最多5个）→ 返回这些步骤的完整内容(detail)。\n"+ - "典型流程：先①/②定位可疑步骤的 step_id，再用③取其完整输出。不含思考(thinking)步骤。", + "典型流程：先①/②定位可疑步骤的 step_id，再用③取其完整输出。不含思考(thinking)步骤。支持直接关联任务的历史 trace；其结果带 source_task_id/inherited=true 且只读。", obj(map[string]any{ "intent_id": idp("意图 id（= work 句柄）"), "q": str("关键字：只返回摘要/完整输出命中它的步骤（可选；与 step_ids 互斥）"), @@ -1111,6 +1484,13 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if id <= 0 { return actool.Errorf("intent_id 必填"), nil } + intentNode, nodeErr := t.ts.GetNodeWithSources(id) + if nodeErr != nil { + return actool.Errorf(nodeErr.Error()), nil + } + if intentNode == nil || intentNode.Kind != db.KindIntent { + return actool.Errorf("intent_id 不属于本任务或其直接关联任务"), nil + } // ③ detail drill-down by step ids (max 5), thinking excluded by the store. if len(a.StepIDs) > 0 { var ids []int64 @@ -1122,31 +1502,47 @@ func (t *ToolSet) getWorkerTrace() actool.CoreTool { if len(ids) > 5 { return actool.Errorf("step_ids 一次最多 5 个（先用不带 step_ids 的调用定位，再分批取）"), nil } - acts, err := t.ts.ActivityByIDs(ids) + acts, err := t.ts.ActivityByIDsWithSources(ids) if err != nil { return actool.Errorf(err.Error()), nil } steps := make([]map[string]any, 0, len(acts)) for i := range acts { - steps = append(steps, map[string]any{ + if acts[i].NodeID == nil || *acts[i].NodeID != id || acts[i].Inherited != intentNode.Inherited || + (acts[i].Inherited && acts[i].SourceTaskID != intentNode.SourceTaskID) { + continue + } + step := map[string]any{ "step_id": acts[i].ID, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "detail": acts[i].Detail, - }) + } + if acts[i].Inherited { + inheritedMap(step, acts[i].SourceTaskID) + } + steps = append(steps, step) } - return jsonResult(map[string]any{"intent_id": id, "steps": steps}) + result := map[string]any{"intent_id": id, "steps": steps} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) } // ①/② summary stream, optionally keyword-filtered; 100-char summaries. var acts []db.Activity var err error if strings.TrimSpace(a.Q) != "" { - acts, err = t.ts.ActivityTraceSearch(&id, a.Q, a.Limit) + acts, err = t.ts.ActivityTraceSearchWithSources(id, a.Q, a.Limit) } else { - acts, err = t.ts.ActivityTrace(id, a.Limit) + acts, err = t.ts.ActivityTraceWithSources(id, a.Limit) } if err != nil { return actool.Errorf(err.Error()), nil } - return jsonResult(map[string]any{"intent_id": id, "steps": traceSteps(acts)}) + result := map[string]any{"intent_id": id, "steps": traceSteps(acts)} + if intentNode.Inherited { + inheritedMap(result, intentNode.SourceTaskID) + } + return jsonResult(result) }) } @@ -1172,7 +1568,7 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { return actool.Errorf("q 必填"), nil } // 排除调用者自身这条意图的步骤（worker 的自有 trace 已在其上下文里）。 - acts, err := t.ts.ActivityTraceSearchExcluding(t.ownerNode, a.Q, a.Limit) + acts, err := t.ts.ActivityTraceSearchAllWithSources(t.ownerNode, a.Q, a.Limit) if err != nil { return actool.Errorf(err.Error()), nil } @@ -1182,11 +1578,15 @@ func (t *ToolSet) searchAllWorkerTraces() actool.CoreTool { if acts[i].NodeID != nil { intent = *acts[i].NodeID } - hits = append(hits, map[string]any{ + hit := map[string]any{ "intent_id": intent, "step_id": acts[i].ID, "worker": acts[i].Worker, "kind": acts[i].Kind, "tool": acts[i].Tool, "is_error": acts[i].IsError, "summary": firstLine(acts[i].Summary, 100), - }) + } + if acts[i].Inherited { + inheritedMap(hit, acts[i].SourceTaskID) + } + hits = append(hits, hit) } return jsonResult(map[string]any{"query": a.Q, "hits": hits}) }) @@ -1217,13 +1617,16 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if limit <= 0 { limit = 50 } - all, err := t.ts.ListByKind(db.KindIntent, 500) + all, err := t.ts.ListByKindWithSources(db.KindIntent, 500) if err != nil { return actool.Errorf(err.Error()), nil } q := strings.ToLower(strings.TrimSpace(a.Q)) out := make([]map[string]any, 0, limit) for _, n := range all { + if n.Inherited && n.State == "running" { + continue + } switch n.State { case "running", "done", "exhausted", "blocked", "stopped": // has run → has a process default: @@ -1235,7 +1638,11 @@ func (t *ToolSet) listWorkerTraces() actool.CoreTool { if q != "" && !strings.Contains(strings.ToLower(summary), q) { continue } - out = append(out, map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State}) + item := map[string]any{"intent_id": n.ID, "summary": summary, "state": n.State} + if n.Inherited { + inheritedMap(item, n.SourceTaskID) + } + out = append(out, item) if len(out) >= limit { break } diff --git a/agent/tools_insert.go b/agent/tools_insert.go index c808988..65f7cf2 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -398,13 +398,12 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { ) } -// listUntestedAssets lets the plan agent pull the CURRENT TASK's in-scope, not-yet- -// tested assets on demand (filter by type, paginated) — so it can decide what to -// test next itself, instead of coverage pushing a backlog list into every prompt. +// listUntestedAssets lets the plan agent pull the current + directly inherited +// scope's not-yet-tested assets on demand (filter by type, paginated). func (t *ToolSet) listUntestedAssets() actool.CoreTool { return readTool( "list_untested_assets", - "查询【本任务】范围内、还没被测过的资产（供你自己判断要不要补测，不代替你决策）。\n"+ + "查询【本任务及直接关联任务】范围内、还没被事实锚点覆盖的资产（关联范围只读，供你自己判断要不要补测，不代替你决策）。\n"+ "可选按资产类型过滤：root_domain/subdomain/service/app/endpoint/ip。\n"+ "分页：page 从 1 起、page_size 默认 10。返回 {assets:[{id,type,label}], total, page, page_size}。仅任务上下文可用。", obj(map[string]any{ @@ -432,7 +431,7 @@ func (t *ToolSet) listUntestedAssets() actool.CoreTool { a.PageSize = 10 } offset := (a.Page - 1) * a.PageSize - assets, total, err := t.as.ListUntestedAssets(t.taskID, t.ts.ID(), strings.TrimSpace(a.Type), a.PageSize, offset) + assets, total, err := t.as.ListUntestedAssetsWithSources(t.taskID, strings.TrimSpace(a.Type), a.PageSize, offset) if err != nil { return actool.Errorf(err.Error()), nil } diff --git a/agent/tools_overview_test.go b/agent/tools_overview_test.go new file mode 100644 index 0000000..8fa6a80 --- /dev/null +++ b/agent/tools_overview_test.go @@ -0,0 +1,31 @@ +package agent + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/Autumn-27/artex/db" +) + +func TestOverviewTextBudgetIsFairAndUTF8Safe(t *testing.T) { + if got := relatedOverviewBudgetForSources(1); got != relatedOverviewMaxTextPerSource { + t.Fatalf("single source budget=%d want=%d", got, relatedOverviewMaxTextPerSource) + } + perSource := relatedOverviewBudgetForSources(db.MaxTaskSourceCount) + if perSource*db.MaxTaskSourceCount > relatedOverviewTotalTextRunes { + t.Fatalf("aggregate budget exceeded: per_source=%d", perSource) + } + + budget := overviewTextBudget{remaining: 5} + got := budget.take(strings.Repeat("中", 10), 20) + if !utf8.ValidString(got) { + t.Fatalf("budget truncation produced invalid UTF-8: %q", got) + } + if utf8.RuneCountInString(got) != 5 || !budget.truncated || budget.remaining != 0 { + t.Fatalf("unexpected truncation: got=%q budget=%+v", got, budget) + } + if tail := budget.take("more", 20); tail != "" { + t.Fatalf("exhausted budget returned more text: %q", tail) + } +} diff --git a/agent/worker.go b/agent/worker.go index 71e1eb9..47e25e9 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -49,6 +49,7 @@ type Worker struct { mem *memory.Store // cross-engagement tradecraft memory (G4) tx *transcript.Store // raw LLM conversation persistence (nil = off) window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum maxTurns int // max agent turns per run (0 = unlimited) // runTimeout is the wall-clock budget for the main exploration of one intent // (0 = unlimited). When it fires, the run is cut and a settlement round is @@ -84,6 +85,15 @@ func NewWorker(prov llm.Provider, model, workDir string, tx *transcript.Store, w return &Worker{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, extraTools: extra} } +func (w *Worker) SetCompactionWindowResolver(fn func() int) { w.windowFn = fn } + +func (w *Worker) compactionWindow() int { + if w.windowFn != nil { + return w.windowFn() + } + return w.window +} + // SetProxy configures the recording proxy address that workers route target // traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through // that MITM proxy. Empty addr disables the hint. @@ -306,8 +316,8 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. // large tool output spills to cmd-output/ with a head + pointer (SDK tool.Capture); // full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。 ToolOutputDir: cmdOutDir(runDir), - Compaction: compactionConfig(w.window), // long tool-heavy runs stay within the window - Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 + Compaction: compactionConfig(w.compactionWindow()), // long tool-heavy runs stay within the window + Todos: actool.NewTodoStore(), // 会话级临时待办（TodoWrite），纯规划用，退出即丢 } if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic) opts.Hooks = hooks diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..7d046cc --- /dev/null +++ b/build.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# ARTEX 单二进制构建脚本：静态导出前端并嵌入 Go 二进制。 +# +# 可选环境变量： +# ARTEX_TARGET_OS=linux 目标系统，默认当前系统 +# ARTEX_TARGET_ARCH=amd64 目标架构，默认当前架构 +# ARTEX_BUILD_VERSION=v0.3.0 构建版本，默认取 git describe +# ARTEX_OUTPUT=/path/to/artex 输出路径 +# ARTEX_SKIP_NPM_CI=1 跳过 npm ci，复用现有依赖 +# ARTEX_GOSUMDB=sum.golang.org Go checksum database +set -euo pipefail + +cd "$(cd "$(dirname "$0")" && pwd)" + +info() { printf '\033[36m[*]\033[0m %s\n' "$*"; } +ok() { printf '\033[32m[+]\033[0m %s\n' "$*"; } +die() { printf '\033[31m[x]\033[0m %s\n' "$*" >&2; exit 1; } + +command -v go >/dev/null 2>&1 || die "未检测到 Go（项目需要 Go 1.26 或更高版本）" +command -v npm >/dev/null 2>&1 || die "未检测到 npm（前端静态构建需要 Node.js/npm）" +command -v rsync >/dev/null 2>&1 || die "未检测到 rsync" + +ARTEX_GOSUMDB="${ARTEX_GOSUMDB:-sum.golang.org}" +ARTEX_TARGET_OS="${ARTEX_TARGET_OS:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOOS)}" +ARTEX_TARGET_ARCH="${ARTEX_TARGET_ARCH:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOARCH)}" + +if [ -z "${ARTEX_BUILD_VERSION:-}" ]; then + if command -v git >/dev/null 2>&1 && git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + ARTEX_BUILD_VERSION="$(git describe --tags --always --dirty)" + else + ARTEX_BUILD_VERSION="dev" + fi +fi + +ARTEX_BINARY_NAME="artex" +[ "$ARTEX_TARGET_OS" = "windows" ] && ARTEX_BINARY_NAME="artex.exe" +ARTEX_OUTPUT="${ARTEX_OUTPUT:-dist/artex-${ARTEX_TARGET_OS}-${ARTEX_TARGET_ARCH}/${ARTEX_BINARY_NAME}}" + +info "构建前端静态资源" +if [ "${ARTEX_SKIP_NPM_CI:-0}" != "1" ]; then + (cd web && npm ci) +fi +(cd web && npm run build:static) + +info "同步前端资源到 server/webui/dist" +mkdir -p server/webui/dist +rsync -a --delete web/out/ server/webui/dist/ + +info "编译 ${ARTEX_TARGET_OS}/${ARTEX_TARGET_ARCH}，版本 ${ARTEX_BUILD_VERSION}" +mkdir -p "$(dirname "$ARTEX_OUTPUT")" +GOSUMDB="$ARTEX_GOSUMDB" \ +CGO_ENABLED=0 \ +GOOS="$ARTEX_TARGET_OS" \ +GOARCH="$ARTEX_TARGET_ARCH" \ +go build \ + -tags embedui \ + -trimpath \ + -ldflags "-s -w -X main.version=${ARTEX_BUILD_VERSION}" \ + -o "$ARTEX_OUTPUT" \ + ./cmd/artex + +ok "编译完成：$ARTEX_OUTPUT" +if command -v file >/dev/null 2>&1; then + file "$ARTEX_OUTPUT" +fi diff --git a/cmd/artex/main.go b/cmd/artex/main.go index b5579dc..cbc5b42 100644 --- a/cmd/artex/main.go +++ b/cmd/artex/main.go @@ -73,14 +73,8 @@ func main() { sigCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - // re-wrap with a named cause so a run cancelled by shutdown says "后端进程正在 - // 关停" in its trace instead of a bare context.Canceled (see agent/cancelcause.go). - ctx, shutdown := context.WithCancelCause(sigCtx) + ctx, shutdown := shutdownContext(sigCtx) defer shutdown(agent.AbortShutdown) - go func() { - <-sigCtx.Done() - shutdown(agent.AbortShutdown) - }() mgr, err := server.NewManager(*dataDir, *proxy) if err != nil { @@ -113,3 +107,18 @@ func main() { defer cancel() _ = httpSrv.Shutdown(shutdownCtx) } + +// shutdownContext deliberately does not derive from signalCtx. If it did, the +// parent's plain context.Canceled could win the race before AbortShutdown was +// attached to the child, losing the diagnostic cause in every running Agent. +func shutdownContext(signalCtx context.Context) (context.Context, context.CancelCauseFunc) { + ctx, shutdown := context.WithCancelCause(context.Background()) + go func() { + select { + case <-signalCtx.Done(): + shutdown(agent.AbortShutdown) + case <-ctx.Done(): + } + }() + return ctx, shutdown +} diff --git a/cmd/artex/main_test.go b/cmd/artex/main_test.go new file mode 100644 index 0000000..149c7c6 --- /dev/null +++ b/cmd/artex/main_test.go @@ -0,0 +1,19 @@ +package main + +import ( + "context" + "testing" + + "github.com/Autumn-27/artex/agent" +) + +func TestShutdownContextPreservesNamedCause(t *testing.T) { + signalCtx, signalCancel := context.WithCancel(context.Background()) + ctx, shutdown := shutdownContext(signalCtx) + defer shutdown(nil) + signalCancel() + <-ctx.Done() + if code, _, _, ok := agent.AbortReason(ctx); !ok || code != "shutdown" { + t.Fatalf("code=%q ok=%v, want shutdown", code, ok) + } +} diff --git a/db/assets.go b/db/assets.go index 583cc87..c1bcb47 100644 --- a/db/assets.go +++ b/db/assets.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "net/url" + "sort" "strings" ) @@ -913,7 +914,7 @@ func (s *AssetStore) CountByType(typ string) (int, error) { } // QueryByCompany returns assets for a company, optionally filtered by type. -// limit <= 0 means no limit — pair it with CountByCompany for server-side paging. +// limit <= 0 means no limit. func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), @@ -940,8 +941,6 @@ FROM assets WHERE company_id = $1` return scanAssets(rows) } -// CountByCompany returns the full number of assets attributed to a company -// (optionally filtered by type), ignoring limit/offset. func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) { q := `SELECT count(*) FROM assets WHERE company_id = $1` args := []any{companyID} @@ -954,8 +953,6 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) { return n, err } -// pageClause appends limit/offset placeholders to args and returns the trailing -// ORDER BY … LIMIT … OFFSET fragment. limit <= 0 means no limit. func pageClause(args *[]any, limit, offset int) string { q := ` ORDER BY last_seen DESC` if limit > 0 { @@ -970,7 +967,6 @@ func pageClause(args *[]any, limit, offset int) string { } // QueryByTask returns assets rows that have a given task_id in task_ids. -// limit <= 0 means no limit — pair it with CountByTask for server-side paging. func (s *AssetStore) QueryByTask(taskID int64, typ string, limit, offset int) ([]*Asset, error) { q := `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), @@ -997,8 +993,6 @@ FROM assets WHERE $1 = ANY(task_ids)` return scanAssets(rows) } -// CountByTask returns the full number of assets attached to a task (optionally -// filtered by type), ignoring limit/offset — lets the UI paginate on the server. func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { q := `SELECT count(*) FROM assets WHERE $1 = ANY(task_ids)` args := []any{taskID} @@ -1011,7 +1005,6 @@ func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { return n, err } -// CountsByTypeForTask returns asset counts per type for a single task. func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) { rows, err := s.db.Query(`SELECT type, COUNT(*) FROM assets WHERE $1 = ANY(task_ids) GROUP BY type`, taskID) if err != nil { @@ -1021,6 +1014,151 @@ func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) { return scanTypeCounts(rows) } +// DeleteByTaskID removes assets owned only by taskID and detaches taskID from +// assets shared with other tasks. Full task deletion uses the coordinated +// transaction in DeleteTaskCascadePrepared; this method remains for callers +// that explicitly manage only asset associations. +func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) { + res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID) + if err != nil { + return 0, err + } + deleted, _ := res.RowsAffected() + if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil { + return deleted, err + } + return deleted, nil +} + +// HostsByTask returns the exact HTTP host candidates attached to a task's +// assets. Domain/IP columns cover root domains, subdomains and non-HTTP +// services; URL covers HTTP services and endpoints. +func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) { + rows, err := s.db.Query(` +SELECT COALESCE(domain,''), COALESCE(ip,''), COALESCE(url,'') +FROM assets WHERE $1 = ANY(task_ids)`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + + hosts := make(map[string]struct{}) + add := func(host string) { + host = strings.TrimSpace(strings.ToLower(host)) + if host != "" { + hosts[host] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + if err := rows.Scan(&domain, &ip, &rawURL); err != nil { + return nil, err + } + add(domain) + add(ip) + if u, err := url.Parse(rawURL); err == nil { + add(u.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(hosts)) + for host := range hosts { + out = append(out, host) + } + sort.Strings(out) + return out, nil +} + +// HostsForTaskDeletion returns hosts that belong to the task being deleted and +// are not referenced by any other live task. Current-task candidates include +// both task_ids ownership and exploration anchors so legacy seeded assets (which +// were anchor-only) are covered. Protection is host-wide: if another live task +// references any asset for a candidate host, that host's global traffic remains. +func (s *AssetStore) HostsForTaskDeletion(taskID, explorationID int64) ([]string, error) { + return hostsForTaskDeletion(s.db, taskID, explorationID) +} + +type rowsQuerier interface { + Query(query string, args ...any) (*sql.Rows, error) +} + +// hostsForTaskDeletion is shared by the read-only AssetStore API and the task +// deletion transaction. Coordinated traffic deletion must call it through the +// transaction path so asset/anchor writes stay locked until the task delete is +// committed. +func hostsForTaskDeletion(q rowsQuerier, taskID, explorationID int64) ([]string, error) { + rows, err := q.Query(` +WITH current_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 +), +other_assets AS ( + SELECT DISTINCT a.id + FROM assets a + WHERE EXISTS ( + SELECT 1 FROM tasks t + WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids) + ) OR 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 + ) +) +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), true +FROM assets a JOIN current_assets c ON c.id=a.id +UNION ALL +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), false +FROM assets a JOIN other_assets o ON o.id=a.id`, taskID, explorationID) + if err != nil { + return nil, err + } + defer rows.Close() + + candidates := make(map[string]struct{}) + protected := make(map[string]struct{}) + add := func(dst map[string]struct{}, raw string) { + raw = strings.TrimSpace(strings.ToLower(raw)) + if raw != "" { + dst[raw] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + var candidate bool + if err := rows.Scan(&domain, &ip, &rawURL, &candidate); err != nil { + return nil, err + } + dst := protected + if candidate { + dst = candidates + } + add(dst, domain) + add(dst, ip) + if u, err := url.Parse(rawURL); err == nil { + add(dst, u.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(candidates)) + for host := range candidates { + if _, shared := protected[host]; !shared { + out = append(out, host) + } + } + sort.Strings(out) + return out, nil +} + const assetSelectCols = `SELECT id, type, company_id, array_to_json(task_ids)::text, COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''), COALESCE(c_segment::text,''), port, @@ -1053,45 +1191,6 @@ func (s *AssetStore) GetByIDs(ids []int64) ([]*Asset, error) { return scanAssets(rows) } -// DeleteByTaskID removes a task's assets: hard-deletes assets owned solely by this -// task (task_ids = {taskID}) and de-associates the rest (drops this task from their -// task_ids, keeping assets shared with other tasks intact). Returns rows deleted. -func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) { - res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID) - if err != nil { - return 0, err - } - n, _ := res.RowsAffected() - if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil { - return n, err - } - return n, nil -} - -// HostsByTask returns the distinct hostnames/IPs the task's assets carry (domain, -// root_domain, ip) — used to purge host-keyed traffic recorded for that task. -func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) { - rows, err := s.db.Query(` -SELECT DISTINCT h FROM ( - SELECT NULLIF(domain,'') AS h FROM assets WHERE $1 = ANY(task_ids) - UNION SELECT NULLIF(root_domain,'') FROM assets WHERE $1 = ANY(task_ids) - UNION SELECT NULLIF(ip,'') FROM assets WHERE $1 = ANY(task_ids) -) q WHERE h IS NOT NULL AND h <> ''`, taskID) - if err != nil { - return nil, err - } - defer rows.Close() - var out []string - for rows.Next() { - var h string - if err := rows.Scan(&h); err != nil { - return nil, err - } - out = append(out, h) - } - return out, rows.Err() -} - // DeleteByCompanyID hard-deletes all assets belonging to a company. Returns rows deleted. func (s *AssetStore) DeleteByCompanyID(companyID int64) (int64, error) { res, err := s.db.Exec(`DELETE FROM assets WHERE company_id = $1`, companyID) @@ -1161,7 +1260,6 @@ func (s *AssetStore) CountsByType() (map[string]int, error) { return scanTypeCounts(rows) } -// scanTypeCounts scans a (type, count) result set into a map. func scanTypeCounts(rows *sql.Rows) (map[string]int, error) { out := map[string]int{} for rows.Next() { diff --git a/db/assets_test.go b/db/assets_test.go index 367f7db..6e92030 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -290,12 +290,12 @@ func TestUpsertHTTPService(t *testing.T) { sc := 200 cl := int64(1024) id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{ - URL: "https://www.httptest.example.com/", - Technologies: []string{"nginx", "vue"}, - StatusCode: &sc, + URL: "https://www.httptest.example.com/", + Technologies: []string{"nginx", "vue"}, + StatusCode: &sc, ContentLength: &cl, - PageTitle: "Test Site", - TaskID: 1, + PageTitle: "Test Site", + TaskID: 1, }) if err != nil { t.Fatal(err) diff --git a/db/config.go b/db/config.go index 67d9edc..8d91487 100644 --- a/db/config.go +++ b/db/config.go @@ -1,8 +1,11 @@ package db import ( + "context" "database/sql" "encoding/json" + "errors" + "fmt" "time" ) @@ -133,9 +136,236 @@ VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$ return p.ID, err } +var ( + ErrActiveLLMProfileDelete = errors.New("cannot delete the active LLM profile; activate another profile first") + ErrLLMProfileReferencesChanged = errors.New("LLM profile references changed while deleting; retry the request") + ErrLLMProfileNotFound = errors.New("LLM profile not found") +) + func (d *DB) DeleteProfile(id int64) error { - _, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) - return err + return d.DeleteProfileContext(context.Background(), id) +} + +// DeleteProfileContext removes a non-default profile while preserving the task +// failover cursor. Reference changes are retried because a task, agent, or +// conversation may start pointing at the profile between the initial scan and +// the profile row lock. A bound prevents a continuously changing workload from +// keeping an HTTP request alive forever. +func (d *DB) DeleteProfileContext(ctx context.Context, id int64) error { + const ( + maxAttempts = 8 + maxDuration = 15 * time.Second + ) + ctx, cancel := context.WithTimeout(ctx, maxDuration) + defer cancel() + for attempt := 0; attempt < maxAttempts; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + retry, err := d.deleteProfile(ctx, id) + if err != nil || !retry { + return err + } + } + return fmt.Errorf("%w: profile %d", ErrLLMProfileReferencesChanged, id) +} + +func (d *DB) deleteProfile(ctx context.Context, id int64) (bool, error) { + tx, err := d.BeginTx(ctx, nil) + if err != nil { + return false, err + } + defer tx.Rollback() + + // Reference setters lock their task/agent/conversation row before an + // llm_profiles row can be locked by a foreign-key check. Keep deletion in the + // same order: the old profile -> child order deadlocked with a setter's child + // -> profile order. Sorting also gives concurrent profile deletions a stable + // order when their chains overlap multiple tasks. + lockedTasks, err := lockProfileReferenceRows(ctx, tx, `SELECT t.id +FROM tasks t +WHERE t.llm_profile_id=$1 + OR t.active_llm_profile_id=$1 + OR EXISTS ( + SELECT 1 FROM task_llm_profiles x + WHERE x.task_id=t.id AND x.profile_id=$1 + ) +ORDER BY t.id +FOR UPDATE OF t`, id) + if err != nil { + return false, err + } + lockedAgents, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM agents +WHERE llm_profile_id=$1 +ORDER BY id +FOR UPDATE`, id) + if err != nil { + return false, err + } + lockedConversations, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM conversations +WHERE llm_profile_id=$1 +ORDER BY id +FOR UPDATE`, id) + if err != nil { + return false, err + } + + var isDefault bool + if err := tx.QueryRowContext(ctx, `SELECT is_default FROM llm_profiles WHERE id=$1 FOR UPDATE`, id).Scan(&isDefault); err != nil { + if err == sql.ErrNoRows { + return false, ErrLLMProfileNotFound + } + return false, err + } + if isDefault { + return false, ErrActiveLLMProfileDelete + } + + type affectedTask struct { + id int64 + position int + wasActive bool + } + // A task may have committed a new reference after the first statement took + // its snapshot but before this transaction acquired the profile lock. The + // profile lock now prevents further references; retry if that committed task + // was not part of the task-first lock set. Never acquire a new task lock while + // holding the profile lock, because that would recreate the inversion. + rows, err := tx.QueryContext(ctx, `SELECT ref_kind, ref_id FROM ( + SELECT 'task'::text AS ref_kind, t.id AS ref_id +FROM tasks t +WHERE t.llm_profile_id=$1 + OR t.active_llm_profile_id=$1 + OR EXISTS ( + SELECT 1 FROM task_llm_profiles x + WHERE x.task_id=t.id AND x.profile_id=$1 + ) + UNION ALL + SELECT 'agent', a.id FROM agents a WHERE a.llm_profile_id=$1 + UNION ALL + SELECT 'conversation', c.id FROM conversations c WHERE c.llm_profile_id=$1 +) refs +ORDER BY ref_kind, ref_id`, id) + if err != nil { + return false, err + } + for rows.Next() { + var ( + kind string + rowID int64 + ) + if err := rows.Scan(&kind, &rowID); err != nil { + rows.Close() + return false, err + } + locked := false + switch kind { + case "task": + _, locked = lockedTasks[rowID] + case "agent": + _, locked = lockedAgents[rowID] + case "conversation": + _, locked = lockedConversations[rowID] + } + if !locked { + rows.Close() + return true, nil + } + } + if err := rows.Err(); err != nil { + rows.Close() + return false, err + } + if err := rows.Close(); err != nil { + return false, err + } + + rows, err = tx.QueryContext(ctx, `SELECT x.task_id, x.position, COALESCE(t.active_llm_profile_id=$1, false) +FROM task_llm_profiles x +JOIN tasks t ON t.id=x.task_id +WHERE x.profile_id=$1 +ORDER BY x.task_id`, id) + if err != nil { + return false, err + } + var affected []affectedTask + for rows.Next() { + var task affectedTask + if err := rows.Scan(&task.id, &task.position, &task.wasActive); err != nil { + rows.Close() + return false, err + } + affected = append(affected, task) + } + if err := rows.Err(); err != nil { + rows.Close() + return false, err + } + if err := rows.Close(); err != nil { + return false, err + } + if _, err := tx.ExecContext(ctx, `DELETE FROM llm_profiles WHERE id=$1`, id); err != nil { + return false, err + } + // Only a successor after the deleted cursor is eligible. If there is none, + // clear the explicit chain so the task falls back to its Agent/global provider; + // profiles before a manually selected cursor must never be revived. + for _, task := range affected { + if !task.wasActive { + if _, err := tx.ExecContext(ctx, `UPDATE tasks SET llm_chain_revision=llm_chain_revision+1 WHERE id=$1`, task.id); err != nil { + return false, err + } + continue + } + var next int64 + err := tx.QueryRowContext(ctx, `SELECT profile_id FROM task_llm_profiles +WHERE task_id=$1 AND position>$2 AND status='ready' +ORDER BY position +LIMIT 1`, task.id, task.position).Scan(&next) + if err != nil && err != sql.ErrNoRows { + return false, err + } + if err == sql.ErrNoRows { + if _, err := tx.ExecContext(ctx, `DELETE FROM task_llm_profiles WHERE task_id=$1`, task.id); err != nil { + return false, err + } + if _, err := tx.ExecContext(ctx, `UPDATE tasks +SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, task.id); err != nil { + return false, err + } + continue + } + if _, err := tx.ExecContext(ctx, `UPDATE tasks +SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1 +WHERE id=$1`, task.id, next); err != nil { + return false, err + } + } + if err := tx.Commit(); err != nil { + return false, err + } + return false, nil +} + +func lockProfileReferenceRows(ctx context.Context, tx *sql.Tx, query string, profileID int64) (map[int64]struct{}, error) { + rows, err := tx.QueryContext(ctx, query, profileID) + if err != nil { + return nil, err + } + defer rows.Close() + locked := make(map[int64]struct{}) + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + locked[id] = struct{}{} + } + if err := rows.Err(); err != nil { + return nil, err + } + return locked, nil } // SetActiveProfile makes one profile the global default (single-default invariant). @@ -148,9 +378,13 @@ func (d *DB) SetActiveProfile(id int64) error { if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil { return err } - if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil { + res, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id) + if err != nil { return err } + if n, _ := res.RowsAffected(); n != 1 { + return ErrLLMProfileNotFound + } return tx.Commit() } @@ -313,8 +547,45 @@ func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error { // the binding (id == nil) so the agent follows the task/conversation pin, else the // global active profile. Precedence at runtime: agent binding → task/conv pin → active. func (d *DB) SetAgentLLMProfile(key string, id *int64) error { - _, err := d.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE key=$2`, id, key) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + // DeleteProfile locks reference rows before the profile row. Keep the same + // order here so a concurrent rebind cannot form a child/profile deadlock. + var agentID int64 + if err := tx.QueryRow(`SELECT id FROM agents WHERE key=$1 FOR UPDATE`, key).Scan(&agentID); err != nil { + if err == sql.ErrNoRows { + // Preserve the previous UPDATE semantics: an unknown key is a no-op. + return tx.Commit() + } + return err + } + if err := lockLLMProfileForReference(tx, id); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE id=$2`, id, agentID); err != nil { + return err + } + return tx.Commit() +} + +// lockLLMProfileForReference makes the profile side of the shared lock-order +// protocol explicit. Callers must already own the referencing child row. +func lockLLMProfileForReference(tx *sql.Tx, profileID *int64) error { + if profileID == nil { + return nil + } + var lockedID int64 + if err := tx.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR KEY SHARE`, *profileID).Scan(&lockedID); err != nil { + if err == sql.ErrNoRows { + return ErrLLMProfileNotFound + } + return err + } + return nil } // SetAgentWebSearch toggles whether an agent uses network search (still gated by diff --git a/db/config_test.go b/db/config_test.go index 6f1743e..16f5d8c 100644 --- a/db/config_test.go +++ b/db/config_test.go @@ -1,7 +1,9 @@ package db import ( + "context" "encoding/json" + "errors" "testing" ) @@ -53,6 +55,15 @@ func TestPoolProfilesOrder(t *testing.T) { } } +func TestDeleteProfileContextHonorsCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var d DB + if err := d.DeleteProfileContext(ctx, 1); !errors.Is(err, context.Canceled) { + t.Fatalf("DeleteProfileContext error=%v, want context cancellation", err) + } +} + func TestConfigStores(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { @@ -69,6 +80,9 @@ func TestConfigStores(t *testing.T) { if err := d.SetActiveProfile(pid); err != nil { t.Fatal(err) } + if err := d.DeleteProfile(pid); !errors.Is(err, ErrActiveLLMProfileDelete) { + t.Fatalf("deleting active profile error=%v, want %v", err, ErrActiveLLMProfileDelete) + } act, err := d.ActiveProfile() if err != nil || act == nil || act.APIKey != "secret123" { t.Fatalf("active profile/key: %+v err=%v", act, err) diff --git a/db/conversation.go b/db/conversation.go index 9e7e97a..2b2ad41 100644 --- a/db/conversation.go +++ b/db/conversation.go @@ -69,19 +69,60 @@ func scanConv(row interface{ Scan(...any) error }) (Conversation, error) { // CreateConversation opens a new chat thread for agentKey with an initial title. // llmProfileID may be nil to use the globally active profile. func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) { - c, err := scanConv(d.QueryRow(` -INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, $3) -RETURNING `+convCols, agentKey, title, llmProfileID)) + tx, err := d.Begin() if err != nil { return nil, err } + defer tx.Rollback() + + // Insert the child without a profile first. The new row is exclusively owned + // by this transaction before it takes a profile lock, matching DeleteProfile's + // child-row -> profile-row protocol. + c, err := scanConv(tx.QueryRow(` +INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, NULL) +RETURNING `+convCols, agentKey, title)) + if err != nil { + return nil, err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return nil, err + } + if llmProfileID != nil { + c, err = scanConv(tx.QueryRow(`UPDATE conversations SET llm_profile_id=$2 +WHERE id=$1 RETURNING `+convCols, c.ID, llmProfileID)) + if err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } return &c, nil } // UpdateConversationProfile sets (or clears) the LLM profile override for a conversation. func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error { - _, err := d.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + var lockedID int64 + if err := tx.QueryRow(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, id).Scan(&lockedID); err != nil { + if err == sql.ErrNoRows { + // Preserve the previous UPDATE semantics: an unknown id is a no-op. + return tx.Commit() + } + return err + } + if err := lockLLMProfileForReference(tx, llmProfileID); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID); err != nil { + return err + } + return tx.Commit() } // ListConversations returns all threads, most-recently-updated first. diff --git a/db/exploration.go b/db/exploration.go index 94f6086..be01753 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -25,34 +25,40 @@ func utf8Clean(s string) string { // Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint. type Node struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - Payload json.RawMessage `json:"payload"` - Priority int `json:"priority"` - State string `json:"state"` - Origin string `json:"origin,omitempty"` - Owner string `json:"owner,omitempty"` - Anchors []int64 `json:"anchors,omitempty"` - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + Kind string `json:"kind"` + Payload json.RawMessage `json:"payload"` + Priority int `json:"priority"` + State string `json:"state"` + Origin string `json:"origin,omitempty"` + Owner string `json:"owner,omitempty"` + BlockedReason string `json:"blocked_reason,omitempty"` + Anchors []int64 `json:"anchors,omitempty"` + CreatedAt time.Time `json:"created_at"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // Activity is one worker execution step (= old task_activity). type Activity struct { - ID int64 `json:"id"` - NodeID *int64 `json:"node_id,omitempty"` - Worker string `json:"worker,omitempty"` - Kind string `json:"kind,omitempty"` - Tool string `json:"tool,omitempty"` - ToolUseID string `json:"tool_use_id,omitempty"` - IsError bool `json:"is_error"` - Summary string `json:"summary,omitempty"` - Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads - CreatedAt time.Time `json:"created_at"` + ID int64 `json:"id"` + NodeID *int64 `json:"node_id,omitempty"` + Worker string `json:"worker,omitempty"` + Kind string `json:"kind,omitempty"` + Tool string `json:"tool,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + IsError bool `json:"is_error"` + Summary string `json:"summary,omitempty"` + Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads + Metadata json.RawMessage `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` // Token usage — set only on the terminal kind='result' record; nil elsewhere. - InputTokens *int `json:"input_tokens,omitempty"` - OutputTokens *int `json:"output_tokens,omitempty"` - CacheReadTokens *int `json:"cache_read_tokens,omitempty"` - CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + InputTokens *int `json:"input_tokens,omitempty"` + OutputTokens *int `json:"output_tokens,omitempty"` + CacheReadTokens *int `json:"cache_read_tokens,omitempty"` + CacheWriteTokens *int `json:"cache_write_tokens,omitempty"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // TokenUsage is a per-worker token aggregate (TokenStatsByWorker). @@ -201,17 +207,17 @@ ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, re // SetNodeState updates any node's state (never deletes). func (s *ExplorationStore) SetNodeState(id int64, state string) error { - _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) + _, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1, blocked_reason=NULL WHERE id=$2 AND exploration_id=$3`, state, id, s.expID) return err } -// SetIntentState updates an intent node's state (done/blocked/exhausted/open). +// SetIntentState updates an intent node's state. // ResetRunningIntents returns any intent left in 'running' back to 'open' so it // is re-claimed. Called on startup: a 'running' intent with no live worker (a // backend restart or crashed worker goroutine left it stuck) would otherwise spin // forever in the UI. Workers re-run an intent from scratch, so reopening is safe. func (s *ExplorationStore) ResetRunningIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) if err != nil { return 0, err @@ -226,7 +232,7 @@ WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) // changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button. func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) { res, err := s.db.Exec(`UPDATE exploration_nodes -SET state='open', completed_at=NULL +SET state='open', completed_at=NULL, blocked_reason=NULL WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state IN ('blocked','exhausted','stopped')`, id, s.expID) if err != nil { @@ -240,7 +246,7 @@ WHERE id=$1 AND exploration_id=$2 AND kind='intent' // (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. func (s *ExplorationStore) ReopenBlockedIntents() (int64, error) { - res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL + res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID) if err != nil { return 0, err @@ -253,17 +259,92 @@ func (s *ExplorationStore) SetIntentState(id int64, state string) error { // terminal states stamp completed_at; reopening (back to open/running) clears it. terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped" _, err := s.db.Exec(`UPDATE exploration_nodes -SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END +SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $4 THEN now() ELSE NULL END WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal) return err } -const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at` +// IntentCleanup summarizes the blackboard records removed when a user cancels +// one worker intent. Global assets and traffic are deliberately not part of this +// operation; exploration anchors disappear only because their owning nodes do. +type IntentCleanup struct { + Intents int64 `json:"intents"` + Facts int64 `json:"facts"` + Findings int64 `json:"findings"` + Activities int64 `json:"activities"` +} + +// CancelIntent removes one local intent and the fact/finding nodes it directly +// yielded. The standalone finding row must be deleted before its node (whose FK +// otherwise uses ON DELETE SET NULL), so the entire cleanup is kept in one DB +// transaction. The caller must first stop a running worker to prevent late writes. +func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) { + var out IntentCleanup + tx, err := s.db.Begin() + if err != nil { + return out, err + } + defer tx.Rollback() + + var state string + if err := tx.QueryRow(`SELECT state FROM exploration_nodes + WHERE id=$1 AND exploration_id=$2 AND kind='intent' FOR UPDATE`, id, s.expID).Scan(&state); err != nil { + if err == sql.ErrNoRows { + return out, fmt.Errorf("intent not found") + } + return out, err + } + if state != "running" && state != "paused" { + return out, fmt.Errorf("intent state %s cannot be cancelled", state) + } + + if err := tx.QueryRow(`SELECT + COUNT(*) FILTER (WHERE n.kind='fact'), + COUNT(*) FILTER (WHERE n.kind='finding') + FROM exploration_edges e + JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.kind IN ('fact','finding')`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil { + return out, err + } + + if _, err := tx.Exec(`DELETE FROM findings f USING exploration_edges e, exploration_nodes n + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.id=e.dst_id AND n.exploration_id=e.exploration_id + AND n.kind='finding' AND f.node_id=n.id`, s.expID, id, RelYields); err != nil { + return out, err + } + res, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id) + if err != nil { + return out, err + } + out.Activities, _ = res.RowsAffected() + + res, err = tx.Exec(`DELETE FROM exploration_nodes + WHERE exploration_id=$1 AND ( + id=$2 OR id IN ( + SELECT e.dst_id FROM exploration_edges e + JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3 + AND n.kind IN ('fact','finding') + ) + )`, s.expID, id, RelYields) + if err != nil { + return out, err + } + if removed, _ := res.RowsAffected(); removed != 1+out.Facts+out.Findings { + return out, fmt.Errorf("intent cleanup changed concurrently") + } + out.Intents = 1 + return out, tx.Commit() +} + +const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at` func scanNode(sc interface{ Scan(...any) error }) (*Node, error) { var n Node var payload []byte - if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil { + if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.BlockedReason, &n.CreatedAt); err != nil { return nil, err } n.Payload = json.RawMessage(payload) @@ -455,6 +536,31 @@ WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields return out, rows.Err() } +// FindingIntentsTerminal is the inherited-history variant: it returns finding +// lineage only when the producing intent has reached an immutable terminal +// state. This keeps a source task's live worker topology private. +func (s *ExplorationStore) FindingIntentsTerminal() (map[int64]int64, error) { + rows, err := s.db.Query(`SELECT e.dst_id, e.src_id + FROM exploration_edges e + JOIN exploration_nodes finding ON finding.id=e.dst_id AND finding.exploration_id=e.exploration_id + JOIN exploration_nodes intent ON intent.id=e.src_id AND intent.exploration_id=e.exploration_id + WHERE e.exploration_id=$1 AND e.rel=$2 AND finding.kind='finding' + AND intent.kind='intent' AND intent.state IN ('done','blocked','exhausted','stopped')`, s.expID, RelYields) + if err != nil { + return nil, err + } + defer rows.Close() + out := map[int64]int64{} + for rows.Next() { + var findingID, intentID int64 + if err := rows.Scan(&findingID, &intentID); err != nil { + return nil, err + } + out[findingID] = intentID + } + return out, rows.Err() +} + // Frontier returns open intents ordered by priority desc then id (FIFO tiebreak). func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) { if limit <= 0 { @@ -516,12 +622,16 @@ func (s *ExplorationStore) Stats() (map[string]int, error) { // AppendActivity records one worker step and returns its id. func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) { + metadata := a.Metadata + if len(metadata) == 0 { + metadata = json.RawMessage(`{}`) + } var id int64 err := s.db.QueryRow(` -INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13) +INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14) RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail), - a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) + metadata, a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id) return id, err } @@ -659,7 +769,7 @@ func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) } var rows *sql.Rows var err error - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` if nodeID != nil { rows, err = s.db.Query(`SELECT `+cols+` FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit) @@ -673,6 +783,43 @@ FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID, defer rows.Close() out := []Activity{} cursor := sinceID + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, sinceID, err + } + if a.ID > cursor { + cursor = a.ID + } + out = append(out, a) + } + return out, cursor, rows.Err() +} + +// ActivityListForTerminalIntent is the inherited, incremental-read boundary. +// The intent state check and activity read happen in one statement so a source +// intent reopened between API-level checks cannot expose its new in-flight trace. +// Model reasoning and accounting rows are never inherited. +func (s *ExplorationStore) ActivityListForTerminalIntent(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + if limit <= 0 { + limit = 300 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND a.id>$3 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $4`, s.expID, nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + defer rows.Close() + out := []Activity{} + cursor := sinceID for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -719,7 +866,7 @@ func (s *ExplorationStore) ActivityPage(f ActivitySessionFilter, before int64, l if limit <= 0 { limit = 200 } - const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` + const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), metadata, created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens` args := []any{s.expID} cond, cargs := f.cond(len(args) + 1) args = append(args, cargs...) @@ -736,6 +883,51 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) } defer rows.Close() desc := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.Metadata, &a.CreatedAt, + &a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil { + return nil, false, err + } + desc = append(desc, a) + } + if err := rows.Err(); err != nil { + return nil, false, err + } + hasMore := len(desc) > limit + if hasMore { + desc = desc[:limit] + } + // reverse the newest-first window into ascending id order for display. + out := make([]Activity, len(desc)) + for i, a := range desc { + out[len(desc)-1-i] = a + } + return out, hasMore, nil +} + +// ActivityPageForTerminalIntent is the inherited session-history boundary. It +// combines ownership, terminal-state and row-kind checks in one query, closing +// the race where a previously completed source intent is reopened before paging. +func (s *ExplorationStore) ActivityPageForTerminalIntent(nodeID, before int64, limit int) ([]Activity, bool, error) { + if limit <= 0 { + limit = 200 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), + COALESCE(a.tool,''), COALESCE(a.tool_use_id,''), a.is_error, COALESCE(a.summary,''), + a.created_at, a.input_tokens, a.output_tokens, a.cache_read_tokens, a.cache_write_tokens + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 + AND n.kind='intent' AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND ($3 <= 0 OR a.id < $3) + ORDER BY a.id DESC LIMIT $4`, s.expID, nodeID, before, limit+1) + if err != nil { + return nil, false, err + } + defer rows.Close() + desc := []Activity{} for rows.Next() { var a Activity if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt, @@ -751,7 +943,6 @@ ORDER BY id DESC LIMIT $%d`, cond, beforeArg, beforeArg, limitArg) if hasMore { desc = desc[:limit] } - // reverse the newest-first window into ascending id order for display. out := make([]Activity, len(desc)) for i, a := range desc { out[len(desc)-1-i] = a @@ -817,6 +1008,26 @@ ORDER BY id LIMIT $3`, s.expID, nodeID, limit) return scanTrace(rows) } +// ActivityTraceForTerminalIntent returns one inherited work trace only while its +// source intent is still terminal at query time. +func (s *ExplorationStore) ActivityTraceForTerminalIntent(nodeID int64, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 500 + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + ORDER BY a.id LIMIT $3`, s.expID, nodeID, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // ActivityTraceSearch finds steps whose summary OR detail matches q // (case-insensitive), returning summaries only. nodeID != nil scopes to one work; // nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no @@ -844,13 +1055,59 @@ AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like, return scanTrace(rows) } +// ActivityTraceSearchForTerminalIntent is the scoped inherited search variant; +// terminal eligibility is evaluated by the same statement that reads the rows. +func (s *ExplorationStore) ActivityTraceSearchForTerminalIntent(nodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.node_id=$2 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $3 OR a.detail ILIKE $3) + ORDER BY a.id LIMIT $4`, s.expID, nodeID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + +// ActivityTraceSearchTerminalIntents searches inherited worker history without +// exposing planner/main rows or work that is still open/running in the source. +func (s *ExplorationStore) ActivityTraceSearchTerminalIntents(q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + like := "%" + q + "%" + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.worker,''), COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.summary,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.kind NOT IN ('thinking','usage') + AND (a.summary ILIKE $2 OR a.detail ILIKE $2) + ORDER BY a.id LIMIT $3`, s.expID, like, limit) + if err != nil { + return nil, err + } + defer rows.Close() + return scanTrace(rows) +} + // AssetRef is a compact exploration node (intent/fact/finding) that references an // asset via an anchor — for the coverage graph's node drawer. type AssetRef struct { - ID int64 `json:"id"` - Kind string `json:"kind"` - State string `json:"state"` - Summary string `json:"summary"` + ID int64 `json:"id"` + Kind string `json:"kind"` + State string `json:"state"` + Summary string `json:"summary"` + SourceTaskID int64 `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } // AssetRefs returns the intents / facts / findings in THIS exploration anchored to @@ -943,6 +1200,42 @@ FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings. return out, rows.Err() } +// ActivityByIDsForTerminalIntents is the inherited-detail read boundary. It +// deliberately excludes node-less planner/main rows, non-intent activities, live +// source work, and thinking. Local callers that need legacy behavior use +// ActivityByIDs instead. +func (s *ExplorationStore) ActivityByIDsForTerminalIntents(ids []int64) ([]Activity, error) { + if len(ids) == 0 { + return nil, nil + } + ph := make([]string, len(ids)) + args := make([]any, len(ids)+1) + args[0] = s.expID + for i, id := range ids { + ph[i] = fmt.Sprintf("$%d", i+2) + args[i+1] = id + } + rows, err := s.db.Query(`SELECT a.id, a.node_id, COALESCE(a.kind,''), COALESCE(a.tool,''), a.is_error, COALESCE(a.detail,'') + FROM activity a + JOIN exploration_nodes n ON n.id=a.node_id AND n.exploration_id=a.exploration_id + WHERE a.exploration_id=$1 AND a.kind NOT IN ('thinking','usage') AND n.kind='intent' + AND n.state IN ('done','blocked','exhausted','stopped') + AND a.id IN (`+strings.Join(ph, ",")+`) ORDER BY a.id`, args...) + if err != nil { + return nil, err + } + defer rows.Close() + out := []Activity{} + for rows.Next() { + var a Activity + if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil { + return nil, err + } + out = append(out, a) + } + return out, rows.Err() +} + // AddStandaloneFinding writes a finding to the standalone findings table, which // persists across task deletion (task_id / node_id become NULL when the task or // exploration node is deleted). taskID and nodeID may be 0 (stored as NULL). diff --git a/db/exploration_sources.go b/db/exploration_sources.go new file mode 100644 index 0000000..3ccbb96 --- /dev/null +++ b/db/exploration_sources.go @@ -0,0 +1,350 @@ +package db + +import ( + "database/sql" + "sort" +) + +// DirectSourceStore binds one directly related task to its exploration store. +// It is intentionally a read-side helper: callers keep using the receiver store +// for every graph mutation, frontier lookup, and intent claim. +type DirectSourceStore struct { + Task TaskSource + Store *ExplorationStore +} + +// DirectSourceStores resolves the live, direct task relations for this +// exploration. It deliberately queries on every call: relations disappear when +// a source task is deleted, and inherited context must not retain stale rows. +// Sources of a source are never expanded. +func (s *ExplorationStore) DirectSourceStores() ([]DirectSourceStore, error) { + rows, err := s.db.Query(` +SELECT source.id, source.exploration_id, source.description, source.goal, source.status +FROM tasks owner +JOIN task_relations relation ON relation.task_id=owner.id +JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL +WHERE owner.exploration_id=$1 AND owner.deleted_at IS NULL +ORDER BY relation.created_at, source.id`, s.expID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []DirectSourceStore{} + for rows.Next() { + var source TaskSource + if err := rows.Scan(&source.TaskID, &source.ExplorationID, &source.Description, &source.Goal, &source.Status); err != nil { + return nil, err + } + out = append(out, DirectSourceStore{Task: source, Store: s.db.Exploration(source.ExplorationID)}) + } + return out, rows.Err() +} + +// TaskID returns the live task bound to this exploration. Explorations created +// directly in tests or maintenance code have no task and return zero. +func (s *ExplorationStore) TaskID() (int64, error) { + var id int64 + err := s.db.QueryRow(`SELECT id FROM tasks WHERE exploration_id=$1 AND deleted_at IS NULL`, s.expID).Scan(&id) + if err != nil { + if err == sql.ErrNoRows { + return 0, nil + } + return 0, err + } + return id, nil +} + +func markInheritedNode(n *Node, taskID int64) *Node { + if n == nil { + return nil + } + n.SourceTaskID = taskID + n.Inherited = true + return n +} + +func markInheritedActivities(in []Activity, taskID int64) []Activity { + for i := range in { + in[i].SourceTaskID = taskID + in[i].Inherited = true + } + return in +} + +func inheritedIntentTerminal(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false + } +} + +// ListByKindWithSources returns local nodes followed by nodes from each direct +// source in relation order. The limit remains per exploration, matching the +// existing ListByKind contract while ensuring one large task cannot hide all +// inherited context from another source. +func (s *ExplorationStore) ListByKindWithSources(kind string, limit int) ([]*Node, error) { + out, err := s.ListByKind(kind, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + nodes, err := source.Store.ListByKind(kind, limit) + if err != nil { + return nil, err + } + for _, node := range nodes { + if kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + out = append(out, markInheritedNode(node, source.Task.TaskID)) + } + } + return out, nil +} + +// GetNodeWithSources reads a node only when it belongs to this exploration or +// one of its direct sources. Inherited nodes are tagged so tool callers can keep +// them read-only and show their provenance. +func (s *ExplorationStore) GetNodeWithSources(id int64) (*Node, error) { + node, err := s.GetNode(id) + if err != nil || node != nil { + return node, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, err = source.Store.GetNode(id) + if err != nil { + return nil, err + } + if node != nil { + if node.Kind == KindIntent && !inheritedIntentTerminal(node.State) { + continue + } + return markInheritedNode(node, source.Task.TaskID), nil + } + } + return nil, nil +} + +// FindingIntentsWithSources combines finding lineage for the current +// exploration and each direct source. Node ids are globally unique. +func (s *ExplorationStore) FindingIntentsWithSources() (map[int64]int64, error) { + out, err := s.FindingIntents() + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + items, err := source.Store.FindingIntentsTerminal() + if err != nil { + return nil, err + } + for findingID, intentID := range items { + out[findingID] = intentID + } + } + return out, nil +} + +// ActivityTraceWithSources returns a work trace when its intent belongs to the +// current exploration or a direct source. It never searches indirect sources. +func (s *ExplorationStore) ActivityTraceWithSources(nodeID int64, limit int) ([]Activity, error) { + acts, err := s.ActivityTrace(nodeID, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceForTerminalIntent(nodeID, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityListWithSources is the source-aware equivalent used by +// get_worker_output. Node ids are global, so the first owning exploration is +// unambiguous even when the work has not emitted any activity yet. +func (s *ExplorationStore) ActivityListWithSources(nodeID, sinceID int64, limit int) ([]Activity, int64, error) { + node, err := s.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node != nil { + return s.ActivityList(&nodeID, sinceID, limit) + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, sinceID, err + } + for _, source := range sources { + node, err = source.Store.GetNode(nodeID) + if err != nil { + return nil, sinceID, err + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, cursor, err := source.Store.ActivityListForTerminalIntent(nodeID, sinceID, limit) + if err != nil { + return nil, sinceID, err + } + return markInheritedActivities(acts, source.Task.TaskID), cursor, nil + } + return []Activity{}, sinceID, nil +} + +// ActivityDetailWithSources keeps the legacy local-task lookup (including local +// thinking rows), while inherited details are restricted to terminal worker +// intents. Source planner/main rows have no node id and must never become part of +// inherited context. +func (s *ExplorationStore) ActivityDetailWithSources(id int64) (string, error) { + detail, err := s.ActivityDetail(id) + if err != nil || detail != "" { + return detail, err + } + acts, err := s.ActivityByIDsWithSources([]int64{id}) + if err != nil || len(acts) == 0 { + return "", err + } + return acts[0].Detail, nil +} + +// ActivityTraceSearchWithSources performs the scoped worker-trace search used +// by get_worker_trace. A node id resolves to at most one exploration because +// exploration node ids are global. +func (s *ExplorationStore) ActivityTraceSearchWithSources(nodeID int64, q string, limit int) ([]Activity, error) { + acts, err := s.ActivityTraceSearch(&nodeID, q, limit) + if err != nil || len(acts) > 0 { + return acts, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + node, nodeErr := source.Store.GetNode(nodeID) + if nodeErr != nil { + return nil, nodeErr + } + if node == nil || node.Kind != KindIntent { + continue + } + acts, err = source.Store.ActivityTraceSearchForTerminalIntent(nodeID, q, limit) + if err != nil { + return nil, err + } + if len(acts) > 0 { + return markInheritedActivities(acts, source.Task.TaskID), nil + } + } + return []Activity{}, nil +} + +// ActivityTraceSearchAllWithSources searches local worker traces plus every +// direct source. The local owner is excluded only from the current exploration; +// inherited traces are immutable historical context. +func (s *ExplorationStore) ActivityTraceSearchAllWithSources(excludeNodeID int64, q string, limit int) ([]Activity, error) { + if limit <= 0 { + limit = 100 + } + out, err := s.ActivityTraceSearchExcluding(excludeNodeID, q, limit) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityTraceSearchTerminalIntents(q, limit) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +// ActivityByIDsWithSources loads local step details with the legacy behavior. For +// direct sources it only returns rows attached to terminal intents, preventing an +// arbitrary global activity id from exposing source planner/main transcripts. +func (s *ExplorationStore) ActivityByIDsWithSources(ids []int64) ([]Activity, error) { + out, err := s.ActivityByIDs(ids) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + acts, err := source.Store.ActivityByIDsForTerminalIntents(ids) + if err != nil { + return nil, err + } + out = append(out, markInheritedActivities(acts, source.Task.TaskID)...) + } + sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID }) + return out, nil +} + +// AssetRefsWithSources returns anchored nodes from this exploration and each +// direct source. Inherited entries retain their owning task id so API/UI callers +// can present them as immutable context. +func (s *ExplorationStore) AssetRefsWithSources(assetID int64) ([]AssetRef, error) { + out, err := s.AssetRefs(assetID) + if err != nil { + return nil, err + } + sources, err := s.DirectSourceStores() + if err != nil { + return nil, err + } + for _, source := range sources { + refs, err := source.Store.AssetRefs(assetID) + if err != nil { + return nil, err + } + for i := range refs { + if refs[i].Kind == KindIntent && !inheritedIntentTerminal(refs[i].State) { + continue + } + refs[i].SourceTaskID = source.Task.TaskID + refs[i].Inherited = true + out = append(out, refs[i]) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].ID > out[j].ID }) + return out, nil +} diff --git a/db/exploration_sources_test.go b/db/exploration_sources_test.go new file mode 100644 index 0000000..7c66fae --- /dev/null +++ b/db/exploration_sources_test.go @@ -0,0 +1,431 @@ +package db + +import ( + "fmt" + "testing" +) + +func TestInheritedActivityReadsRequireTerminalIntent(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("source activity boundary", "source activity boundary") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) }) + store := d.Exploration(expID) + intentID, err := store.AddIntent(map[string]any{"summary": "source work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + appendStep := func(kind, summary string) int64 { + t.Helper() + id, appendErr := store.AppendActivity(Activity{ + NodeID: &intentID, Worker: "worker", Kind: kind, Summary: summary, Detail: summary + " detail", + }) + if appendErr != nil { + t.Fatal(appendErr) + } + return id + } + textID := appendStep("text", "shared text") + thinkingID := appendStep("thinking", "private reasoning") + usageID := appendStep("usage", "private accounting") + resultID := appendStep("result", "shared result") + if err := store.SetIntentState(intentID, "done"); err != nil { + t.Fatal(err) + } + + page, more, err := store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || more || len(page) != 2 || page[0].ID != textID || page[1].ID != resultID { + t.Fatalf("terminal page boundary: page=%+v more=%v err=%v", page, more, err) + } + list, cursor, err := store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 2 || cursor != resultID { + t.Fatalf("terminal list boundary: list=%+v cursor=%d err=%v", list, cursor, err) + } + trace, err := store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 2 { + t.Fatalf("terminal trace boundary: trace=%+v err=%v", trace, err) + } + hits, err := store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 2 { + t.Fatalf("terminal search boundary: hits=%+v err=%v", hits, err) + } + details, err := store.ActivityByIDsForTerminalIntents([]int64{textID, thinkingID, usageID, resultID}) + if err != nil || len(details) != 2 || details[0].ID != textID || details[1].ID != resultID { + t.Fatalf("terminal detail boundary: details=%+v err=%v", details, err) + } + + // Reopening a source intent must close every inherited activity read even if + // callers still hold a stale terminal-state snapshot. + if err := store.SetIntentState(intentID, "running"); err != nil { + t.Fatal(err) + } + page, _, err = store.ActivityPageForTerminalIntent(intentID, 0, 10) + if err != nil || len(page) != 0 { + t.Fatalf("reopened intent leaked through page: page=%+v err=%v", page, err) + } + list, _, err = store.ActivityListForTerminalIntent(intentID, 0, 10) + if err != nil || len(list) != 0 { + t.Fatalf("reopened intent leaked through list: list=%+v err=%v", list, err) + } + trace, err = store.ActivityTraceForTerminalIntent(intentID, 10) + if err != nil || len(trace) != 0 { + t.Fatalf("reopened intent leaked through trace: trace=%+v err=%v", trace, err) + } + hits, err = store.ActivityTraceSearchForTerminalIntent(intentID, "shared", 10) + if err != nil || len(hits) != 0 { + t.Fatalf("reopened intent leaked through search: hits=%+v err=%v", hits, err) + } + details, err = store.ActivityByIDsForTerminalIntents([]int64{textID, resultID}) + if err != nil || len(details) != 0 { + t.Fatalf("reopened intent leaked through detail: details=%+v err=%v", details, err) + } +} + +func TestExplorationDirectSourceReadView(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand source", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("direct source", "source goal", TaskCreateOptions{ + SourceTaskIDs: []int64{grand.ID}, + }) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current", "current goal", TaskCreateOptions{ + SourceTaskIDs: []int64{source.ID}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + }) + + grandStore := d.Exploration(grand.ExplorationID) + sourceStore := d.Exploration(source.ExplorationID) + currentStore := d.Exploration(current.ExplorationID) + + grandFact, err := grandStore.AddNode(KindFact, map[string]any{"summary": "grand-only fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFact, err := sourceStore.AddNode(KindFact, map[string]any{"summary": "direct source fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceFinding, err := sourceStore.AddNode(KindFinding, map[string]any{"summary": "direct source finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work"}, 9, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.Link(sourceIntent, RelYields, sourceFinding); err != nil { + t.Fatal(err) + } + currentIntent, err := currentStore.AddIntent(map[string]any{"summary": "current work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + currentFact, err := currentStore.AddNode(KindFact, map[string]any{"summary": "current fact"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + + sources, err := currentStore.DirectSourceStores() + if err != nil { + t.Fatal(err) + } + if len(sources) != 1 || sources[0].Task.TaskID != source.ID { + t.Fatalf("want only direct source %d, got %+v", source.ID, sources) + } + + frontier, err := currentStore.Frontier(20) + if err != nil { + t.Fatal(err) + } + if len(frontier) != 1 || frontier[0].ID != currentIntent { + t.Fatalf("source intent leaked into local frontier: %+v", frontier) + } + if claimed, err := currentStore.ClaimIntent(sourceIntent, "wrong-task-worker"); err != nil || claimed { + t.Fatalf("source intent must not be claimable: claimed=%v err=%v", claimed, err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got != nil { + t.Fatalf("live source intent must not be readable as inherited history: node=%+v err=%v", got, err) + } + if got, err := currentStore.ListByKindWithSources(KindIntent, 100); err != nil || len(got) != 1 || got[0].ID != currentIntent { + t.Fatalf("live source intent leaked through inherited list: nodes=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != 0 { + t.Fatalf("live source intent leaked through finding lineage: lineage=%+v err=%v", got, err) + } + + facts, err := currentStore.ListByKindWithSources(KindFact, 100) + if err != nil { + t.Fatal(err) + } + seenCurrent, seenSource, seenGrand := false, false, false + for _, fact := range facts { + switch fact.ID { + case currentFact: + seenCurrent = !fact.Inherited && fact.SourceTaskID == 0 + case sourceFact: + seenSource = fact.Inherited && fact.SourceTaskID == source.ID + case grandFact: + seenGrand = true + } + } + if !seenCurrent || !seenSource || seenGrand { + t.Fatalf("unexpected direct-source fact view: current=%v source=%v grand=%v", seenCurrent, seenSource, seenGrand) + } + + gotSource, err := currentStore.GetNodeWithSources(sourceFact) + if err != nil || gotSource == nil || !gotSource.Inherited || gotSource.SourceTaskID != source.ID { + t.Fatalf("source node lookup: node=%+v err=%v", gotSource, err) + } + gotGrand, err := currentStore.GetNodeWithSources(grandFact) + if err != nil || gotGrand != nil { + t.Fatalf("indirect source must be invisible: node=%+v err=%v", gotGrand, err) + } + + // Existing write methods remain bound to currentStore.expID. Passing an + // inherited id is a no-op and cannot mutate the source blackboard. + if err := currentStore.SetNodeState(sourceFact, "dismissed"); err != nil { + t.Fatal(err) + } + unchanged, err := sourceStore.GetNode(sourceFact) + if err != nil || unchanged == nil || unchanged.State != "confirmed" { + t.Fatalf("inherited fact was mutated: node=%+v err=%v", unchanged, err) + } + + if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil { + t.Fatal(err) + } + if got, err := currentStore.GetNodeWithSources(sourceIntent); err != nil || got == nil || !got.Inherited { + t.Fatalf("terminal source intent should become readable history: node=%+v err=%v", got, err) + } + if got, err := currentStore.FindingIntentsWithSources(); err != nil || got[sourceFinding] != sourceIntent { + t.Fatalf("terminal source finding lineage missing: lineage=%+v err=%v", got, err) + } + sourcePlannerStep, err := sourceStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "private source planner", Detail: "private source planner transcript", + }) + if err != nil { + t.Fatal(err) + } + openSourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "source work still running"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := sourceStore.SetIntentState(openSourceIntent, "running"); err != nil { + t.Fatal(err) + } + openSourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &openSourceIntent, Worker: "source-live-worker", Kind: "text", + Summary: "private live source work", Detail: "private live source transcript", + }) + if err != nil { + t.Fatal(err) + } + sourceStep, err := sourceStore.AppendActivity(Activity{ + NodeID: &sourceIntent, Worker: "source-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "source-secret response", Detail: "source-secret full detail", + }) + if err != nil { + t.Fatal(err) + } + currentPlannerStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "text", Summary: "current planner", Detail: "current planner detail", + }) + if err != nil { + t.Fatal(err) + } + currentThinkingStep, err := currentStore.AppendActivity(Activity{ + Worker: "planner", Kind: "thinking", Summary: "current thinking", Detail: "current thinking detail", + }) + if err != nil { + t.Fatal(err) + } + grandStep, err := grandStore.AppendActivity(Activity{ + NodeID: &grandFact, Worker: "grand-worker", Kind: "tool_result", Tool: "HTTP", + Summary: "grand-only response", Detail: "grand-only detail", + }) + if err != nil { + t.Fatal(err) + } + + trace, err := currentStore.ActivityTraceWithSources(sourceIntent, 20) + if err != nil || len(trace) != 1 || trace[0].ID != sourceStep || !trace[0].Inherited || trace[0].SourceTaskID != source.ID { + t.Fatalf("source trace lookup: trace=%+v err=%v", trace, err) + } + hits, err := currentStore.ActivityTraceSearchAllWithSources(0, "source-secret", 20) + if err != nil || len(hits) != 1 || hits[0].ID != sourceStep || !hits[0].Inherited { + t.Fatalf("source trace search: hits=%+v err=%v", hits, err) + } + grandHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "grand-only", 20) + if err != nil || len(grandHits) != 0 { + t.Fatalf("indirect source trace leaked: hits=%+v err=%v", grandHits, err) + } + liveHits, err := currentStore.ActivityTraceSearchAllWithSources(0, "private live source", 20) + if err != nil || len(liveHits) != 0 { + t.Fatalf("live source trace leaked: hits=%+v err=%v", liveHits, err) + } + if liveTrace, err := currentStore.ActivityTraceWithSources(openSourceIntent, 20); err != nil || len(liveTrace) != 0 { + t.Fatalf("live inherited intent trace leaked: trace=%+v err=%v", liveTrace, err) + } + details, err := currentStore.ActivityByIDsWithSources([]int64{ + sourcePlannerStep, openSourceStep, sourceStep, currentPlannerStep, grandStep, + }) + if err != nil || len(details) != 2 || details[0].ID != sourceStep || !details[0].Inherited || details[1].ID != currentPlannerStep || details[1].Inherited { + t.Fatalf("source trace detail allow-list: details=%+v err=%v", details, err) + } + if detail, err := currentStore.ActivityDetailWithSources(sourcePlannerStep); err != nil || detail != "" { + t.Fatalf("source planner transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(openSourceStep); err != nil || detail != "" { + t.Fatalf("live source intent transcript leaked: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentPlannerStep); err != nil || detail != "current planner detail" { + t.Fatalf("local activity detail behavior changed: detail=%q err=%v", detail, err) + } + if detail, err := currentStore.ActivityDetailWithSources(currentThinkingStep); err != nil || detail != "current thinking detail" { + t.Fatalf("local thinking detail behavior changed: detail=%q err=%v", detail, err) + } +} + +func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + grand, err := d.CreateTask("grand assets", "grand goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + source, err := d.CreateTaskWithOptions("source assets", "source goal", TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}}) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current assets", "current goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + sourceTestedDomain := fmt.Sprintf("context-tested-%d.invalid", source.ID) + sourceUntestedDomain := fmt.Sprintf("context-untested-%d.invalid", source.ID) + sourceAnchoredOnlyDomain := fmt.Sprintf("context-anchor-only-%d.invalid", source.ID) + grandDomain := fmt.Sprintf("context-grand-%d.invalid", grand.ID) + sourceTestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceTestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceUntestedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceUntestedDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + sourceAnchoredOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sourceAnchoredOnlyDomain, TaskID: source.ID}) + if err != nil { + t.Fatal(err) + } + grandID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: grandDomain, TaskID: grand.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(current.ID) + _ = d.DeleteTask(source.ID) + _ = d.DeleteTask(grand.ID) + _, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID}) + }) + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test"); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "tested source asset"}, 0, "confirmed", "worker", []int64{sourceTestedID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "anchored-only source asset"}, 0, "confirmed", "worker", []int64{sourceAnchoredOnlyID}); err != nil { + t.Fatal(err) + } + if _, err := d.Exploration(grand.ExplorationID).AddNode(KindFact, map[string]any{"summary": "indirect tested asset"}, 0, "confirmed", "worker", []int64{grandID}); err != nil { + t.Fatal(err) + } + + scopes, err := assets.ListTaskScopeWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if len(scopes) != 2 || scopes[0].TaskID != source.ID || scopes[1].TaskID != source.ID { + t.Fatalf("direct source scopes only: %+v", scopes) + } + cov, err := assets.TaskCoverageWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + if cov.Denominator != 3 || cov.Tested != 2 { + t.Fatalf("combined coverage: %+v", cov) + } + untested, total, err := assets.ListUntestedAssetsWithSources(current.ID, "", 10, 0) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(untested) != 1 || untested[0].ID != sourceUntestedID { + t.Fatalf("combined untested assets: total=%d assets=%+v", total, untested) + } + hosts, err := assets.HostsByTaskWithSources(current.ID) + if err != nil { + t.Fatal(err) + } + hostSet := map[string]bool{} + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[sourceTestedDomain] || !hostSet[sourceUntestedDomain] || !hostSet[sourceAnchoredOnlyDomain] || hostSet[grandDomain] { + t.Fatalf("direct source hosts only: %v", hosts) + } + refs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(sourceAnchoredOnlyID) + if err != nil || len(refs) != 1 || !refs[0].Inherited || refs[0].SourceTaskID != source.ID || refs[0].Kind != KindFact { + t.Fatalf("source asset refs provenance: refs=%+v err=%v", refs, err) + } + grandRefs, err := d.Exploration(current.ExplorationID).AssetRefsWithSources(grandID) + if err != nil || len(grandRefs) != 0 { + t.Fatalf("indirect source asset refs leaked: refs=%+v err=%v", grandRefs, err) + } + graph, err := assets.BuildCoverageGraph(current.ID, current.ExplorationID) + if err != nil { + t.Fatal(err) + } + seenAnchoredOnly := false + for _, node := range graph.Nodes { + if node.AssetID == sourceAnchoredOnlyID { + seenAnchoredOnly = node.Tested && node.InScope + } + } + if !seenAnchoredOnly { + t.Fatalf("anchored-only source asset missing from coverage graph: %+v", graph.Nodes) + } +} diff --git a/db/exploration_test.go b/db/exploration_test.go index 1c31d9f..97d9f62 100644 --- a/db/exploration_test.go +++ b/db/exploration_test.go @@ -1,6 +1,9 @@ package db -import "testing" +import ( + "fmt" + "testing" +) func TestExplorationFlow(t *testing.T) { d, err := Open(testDSN(t)) @@ -112,3 +115,134 @@ func TestExplorationFlow(t *testing.T) { t.Fatalf("stats: %+v", st) } } + +func TestIntentPauseResumeAndCancelCleanup(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + expID, err := d.CreateExploration("test", "worker control cleanup") + if err != nil { + t.Fatal(err) + } + defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) + es := d.Exploration(expID) + + assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("cancel-intent-%d.invalid", expID), + }) + if err != nil { + t.Fatal(err) + } + defer deleteAsset(d, assetID) + + otherIntent, err := es.AddIntent(map[string]any{"summary": "keep intent"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + intentID, err := es.AddIntent(map[string]any{"summary": "cancel intent"}, 10, nil, "planner") + if err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-control-test"); err != nil || !claimed { + t.Fatalf("initial claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + frontier, err := es.Frontier(10) + if err != nil { + t.Fatal(err) + } + for _, node := range frontier { + if node.ID == intentID { + t.Fatalf("paused intent %d must not enter frontier", intentID) + } + } + if claimed, err := es.ClaimIntent(intentID, "worker-while-paused"); err != nil || claimed { + t.Fatalf("paused intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "open"); err != nil { + t.Fatal(err) + } + if claimed, err := es.ClaimIntent(intentID, "worker-after-resume"); err != nil || !claimed { + t.Fatalf("resumed intent claim: claimed=%v err=%v", claimed, err) + } + if err := es.SetIntentState(intentID, "paused"); err != nil { + t.Fatal(err) + } + + directFact, err := es.AddNode("fact", map[string]any{"summary": "remove fact"}, 0, "confirmed", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + directFinding, err := es.AddNode("finding", map[string]any{"summary": "remove finding"}, 0, "confirmed", "worker", nil) + if err != nil { + t.Fatal(err) + } + keptFact, err := es.AddNode("fact", map[string]any{"summary": "keep fact"}, 0, "confirmed", "other-worker", nil) + if err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFact); err != nil { + t.Fatal(err) + } + if err := es.Link(intentID, RelYields, directFinding); err != nil { + t.Fatal(err) + } + findingRowID, err := es.AddStandaloneFinding(0, directFinding, "test", "cancelled finding", SeverityHigh, "summary", "evidence", "worker", []int64{assetID}) + if err != nil { + t.Fatal(err) + } + activityID, err := es.AppendActivity(Activity{NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "remove activity"}) + if err != nil { + t.Fatal(err) + } + keptActivityID, err := es.AppendActivity(Activity{NodeID: &otherIntent, Worker: "other-worker", Kind: "result", Summary: "keep activity"}) + if err != nil { + t.Fatal(err) + } + + cleanup, err := es.CancelIntent(intentID) + if err != nil { + t.Fatal(err) + } + if cleanup.Intents != 1 || cleanup.Facts != 1 || cleanup.Findings != 1 || cleanup.Activities != 1 { + t.Fatalf("unexpected cleanup counts: %+v", cleanup) + } + for _, nodeID := range []int64{intentID, directFact, directFinding} { + node, err := es.GetNode(nodeID) + if err != nil { + t.Fatal(err) + } + if node != nil { + t.Fatalf("node %d survived intent cancellation", nodeID) + } + } + for _, nodeID := range []int64{otherIntent, keptFact} { + node, err := es.GetNode(nodeID) + if err != nil || node == nil { + t.Fatalf("unrelated node %d removed: node=%v err=%v", nodeID, node, err) + } + } + + assertCount := func(query string, want int, args ...any) { + t.Helper() + var got int + if err := d.QueryRow(query, args...).Scan(&got); err != nil { + t.Fatal(err) + } + if got != want { + t.Fatalf("query count=%d, want %d: %s", got, want, query) + } + } + assertCount(`SELECT COUNT(*) FROM findings WHERE id=$1`, 0, findingRowID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 0, activityID) + assertCount(`SELECT COUNT(*) FROM activity WHERE id=$1`, 1, keptActivityID) + assertCount(`SELECT COUNT(*) FROM exploration_edges WHERE exploration_id=$1 AND (src_id=$2 OR dst_id=$2)`, 0, expID, intentID) + assertCount(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1`, 0, directFact) + assertCount(`SELECT COUNT(*) FROM assets WHERE id=$1`, 1, assetID) +} diff --git a/db/findings.go b/db/findings.go index 6fc64db..9dfe6e2 100644 --- a/db/findings.go +++ b/db/findings.go @@ -9,7 +9,8 @@ import ( "time" ) -// DBFinding is a row in the standalone findings table (persists across task deletion). +// DBFinding is a row in the standalone findings table. It persists across task +// deletion unless the caller explicitly requests related finding cleanup. type DBFinding struct { ID int64 TaskID *int64 diff --git a/db/schema.sql b/db/schema.sql index 32ba6f9..90779f7 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -129,19 +129,40 @@ CREATE TABLE IF NOT EXISTS exploration_nodes ( state TEXT NOT NULL DEFAULT 'open', origin TEXT, owner TEXT, + blocked_reason TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), completed_at TIMESTAMPTZ, CONSTRAINT ck_node_kind CHECK (kind IN ('begin','goal','intent','fact','finding','hint')), CONSTRAINT ck_node_state CHECK ( (kind='begin' AND state IN ('open')) OR - (kind='intent' AND state IN ('open','running','done','blocked','exhausted','stopped')) OR + (kind='intent' AND state IN ('open','running','paused','done','blocked','exhausted','stopped')) OR (kind='goal' AND state IN ('open','met','abandoned')) OR (kind='fact' AND state IN ('confirmed','dismissed','origin')) OR (kind='finding' AND state IN ('confirmed','dismissed')) OR (kind='hint' AND state IN ('active','consumed')) ) ); +ALTER TABLE exploration_nodes ADD COLUMN IF NOT EXISTS blocked_reason TEXT; +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conrelid='exploration_nodes'::regclass + AND conname='ck_node_state' + AND pg_get_constraintdef(oid) NOT LIKE '%paused%' + ) THEN + ALTER TABLE exploration_nodes DROP CONSTRAINT ck_node_state; + ALTER TABLE exploration_nodes ADD CONSTRAINT ck_node_state CHECK ( + (kind='begin' AND state IN ('open')) OR + (kind='intent' AND state IN ('open','running','paused','done','blocked','exhausted','stopped')) OR + (kind='goal' AND state IN ('open','met','abandoned')) OR + (kind='fact' AND state IN ('confirmed','dismissed','origin')) OR + (kind='finding' AND state IN ('confirmed','dismissed')) OR + (kind='hint' AND state IN ('active','consumed')) + ); + END IF; +END $$; CREATE INDEX IF NOT EXISTS idx_expnodes_part ON exploration_nodes(exploration_id, kind); CREATE INDEX IF NOT EXISTS idx_expnodes_frontier ON exploration_nodes(exploration_id, priority DESC) WHERE kind='intent' AND state='open'; @@ -179,12 +200,14 @@ CREATE TABLE IF NOT EXISTS activity ( is_error BOOLEAN NOT NULL DEFAULT false, summary TEXT, detail TEXT, + metadata JSONB NOT NULL DEFAULT '{}', input_tokens INTEGER, output_tokens INTEGER, cache_read_tokens INTEGER, cache_write_tokens INTEGER, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +ALTER TABLE activity ADD COLUMN IF NOT EXISTS metadata JSONB NOT NULL DEFAULT '{}'; CREATE INDEX IF NOT EXISTS idx_act_node ON activity(exploration_id, node_id, id); CREATE INDEX IF NOT EXISTS idx_act_since ON activity(exploration_id, id); -- Main/Plan history pages filter by worker (both carry NULL node_id, so idx_act_node @@ -256,6 +279,8 @@ CREATE TABLE IF NOT EXISTS tasks ( paused BOOLEAN NOT NULL DEFAULT false, queued BOOLEAN NOT NULL DEFAULT false, llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, + active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, + llm_chain_revision BIGINT NOT NULL DEFAULT 0, company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL, parent_ref TEXT, timeout_seconds INTEGER NOT NULL DEFAULT 0, @@ -276,6 +301,57 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300; -- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。 ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_chain_revision BIGINT NOT NULL DEFAULT 0; + +-- Direct, read-only task context inheritance. Relations are intentionally not +-- recursive: a task sees only the source tasks explicitly chosen at creation. +CREATE TABLE IF NOT EXISTS task_relations ( + task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + source_task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (task_id, source_task_id), + CONSTRAINT ck_task_relation_not_self CHECK (task_id <> source_task_id) +); +CREATE INDEX IF NOT EXISTS idx_task_relations_source ON task_relations(source_task_id); + +-- Ordered task-level LLM failover chain. A quota-exhausted entry is skipped +-- until the user saves/resets the chain, which clears all failure state. +CREATE TABLE IF NOT EXISTS task_llm_profiles ( + task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + profile_id BIGINT NOT NULL REFERENCES llm_profiles(id) ON DELETE CASCADE, + position INTEGER NOT NULL CHECK (position >= 0), + status TEXT NOT NULL DEFAULT 'ready' + CHECK (status IN ('ready','quota_exhausted')), + last_error TEXT, + exhausted_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (task_id, profile_id), + UNIQUE (task_id, position) +); +CREATE INDEX IF NOT EXISTS idx_task_llm_profiles_order ON task_llm_profiles(task_id, position); +CREATE INDEX IF NOT EXISTS idx_task_llm_profiles_profile ON task_llm_profiles(profile_id, task_id); +CREATE INDEX IF NOT EXISTS idx_tasks_llm_profile ON tasks(llm_profile_id) WHERE llm_profile_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_tasks_active_llm_profile ON tasks(active_llm_profile_id) WHERE active_llm_profile_id IS NOT NULL; +DROP TRIGGER IF EXISTS trg_task_llm_profiles_upd ON task_llm_profiles; +CREATE TRIGGER trg_task_llm_profiles_upd BEFORE UPDATE ON task_llm_profiles + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + +-- One-time-compatible backfill: old pinned tasks become one-entry chains. A user +-- can still clear the chain later because the update path also clears the legacy +-- llm_profile_id column, preventing this block from re-adding it on restart. +INSERT INTO task_llm_profiles(task_id, profile_id, position) +SELECT t.id, t.llm_profile_id, 0 +FROM tasks t +WHERE t.llm_profile_id IS NOT NULL + AND NOT EXISTS (SELECT 1 FROM task_llm_profiles x WHERE x.task_id=t.id) +ON CONFLICT DO NOTHING; +UPDATE tasks t +SET active_llm_profile_id = t.llm_profile_id +WHERE t.active_llm_profile_id IS NULL + AND t.llm_profile_id IS NOT NULL + AND EXISTS (SELECT 1 FROM task_llm_profiles x WHERE x.task_id=t.id AND x.profile_id=t.llm_profile_id); -- 任务测试范围（资产覆盖度的分母 + 授权边界）。 -- 自动填(source='auto')：insertAssets 顶层按 worker 显式插入的资产类型加保守范围 @@ -335,6 +411,7 @@ ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_merge_mode TEXT NOT NUL ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_max_parallel INTEGER NOT NULL DEFAULT 5; -- per-agent LLM 绑定(agent 级默认模型):列自初版即在上方 CREATE 中,此 ALTER 仅为极旧库兜底(幂等)。 ALTER TABLE agents ADD COLUMN IF NOT EXISTS llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL; +CREATE INDEX IF NOT EXISTS idx_agents_llm_profile ON agents(llm_profile_id) WHERE llm_profile_id IS NOT NULL; DROP TRIGGER IF EXISTS trg_agents_upd ON agents; CREATE TRIGGER trg_agents_upd BEFORE UPDATE ON agents FOR EACH ROW EXECUTE FUNCTION set_updated_at(); @@ -454,6 +531,7 @@ CREATE TABLE IF NOT EXISTS conversations ( created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL; DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations; CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations FOR EACH ROW EXECUTE FUNCTION set_updated_at(); diff --git a/db/task_assets_context.go b/db/task_assets_context.go new file mode 100644 index 0000000..722cb36 --- /dev/null +++ b/db/task_assets_context.go @@ -0,0 +1,215 @@ +package db + +import ( + "net/url" + "sort" + "strconv" + "strings" +) + +// directTaskContextCTE is the current task plus exactly its explicitly related +// source tasks. It is deliberately non-recursive. +const directTaskContextCTE = ` +context_tasks AS ( + SELECT t.id AS task_id, t.exploration_id + FROM tasks t + WHERE t.id=$1 AND t.deleted_at IS NULL + UNION ALL + SELECT source.id, source.exploration_id + FROM task_relations relation + JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL + WHERE relation.task_id=$1 +)` + +const contextCoverageCTE = directTaskContextCTE + `, +target AS ( + SELECT DISTINCT a.id, a.type, + COALESCE(a.url, a.domain, a.ip, a.app_name, a.root_domain, '') AS label + FROM assets a + JOIN task_scope ts ON ( + (ts.kind='company' AND a.company_id = ts.company_id) + OR (ts.kind='root_domain' AND a.root_domain = ts.domain) + OR (ts.kind='subdomain' AND a.domain = ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + ) + JOIN context_tasks ctx ON ctx.task_id=ts.task_id + UNION + SELECT a.id, a.type, + COALESCE(a.url, a.domain, a.ip, a.app_name, a.root_domain, '') AS label + FROM assets a + JOIN exploration_anchors ea ON ea.asset_id=a.id + JOIN exploration_nodes en ON en.id=ea.node_id + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id +), +tested AS ( + SELECT DISTINCT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes en ON en.id=ea.node_id AND en.kind='fact' + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id +)` + +// ListTaskScopeWithSources returns the current task's scope followed by the +// scopes of its direct source tasks. TaskScope.TaskID preserves provenance. +func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) { + rows, err := s.db.Query(`WITH `+directTaskContextCTE+` +SELECT ts.id, ts.task_id, ts.kind, COALESCE(ts.company_id,0), COALESCE(ts.domain,''), + COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'') +FROM task_scope ts +JOIN context_tasks ctx ON ctx.task_id=ts.task_id +ORDER BY CASE WHEN ts.task_id=$1 THEN 0 ELSE 1 END, ts.id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []TaskScope{} + for rows.Next() { + var scope TaskScope + var companyID int64 + if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.Domain, &scope.Net, &scope.Source, &scope.Reason); err != nil { + return nil, err + } + if companyID > 0 { + scope.CompanyID = &companyID + } + out = append(out, scope) + } + return out, rows.Err() +} + +// TaskCoverageWithSources computes one coverage view over the union of the +// current task and its direct sources: source scopes and anchored assets extend +// the denominator, while fact anchors count as tested. No row is copied. +func (s *AssetStore) TaskCoverageWithSources(taskID int64) (*Coverage, error) { + cov := &Coverage{ByType: []CoverageByType{}} + _ = s.db.QueryRow(`WITH `+directTaskContextCTE+` +SELECT count(*) FROM task_scope ts JOIN context_tasks ctx ON ctx.task_id=ts.task_id`, taskID).Scan(&cov.ScopeRows) + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` +SELECT target.type, count(*) AS total, + count(*) FILTER (WHERE target.id IN (SELECT asset_id FROM tested)) AS tested +FROM target GROUP BY target.type ORDER BY target.type`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var byType CoverageByType + if err := rows.Scan(&byType.Type, &byType.Total, &byType.Tested); err != nil { + return nil, err + } + cov.ByType = append(cov.ByType, byType) + cov.Denominator += byType.Total + cov.Tested += byType.Tested + } + if err := rows.Err(); err != nil { + return nil, err + } + if cov.Denominator > 0 { + pct := float64(cov.Tested) / float64(cov.Denominator) + cov.Pct = &pct + } + return cov, nil +} + +// ListUntestedAssetsWithSources is the direct-source-aware backlog query used +// by inherited tasks. Source fact anchors remove assets from the backlog. +func (s *AssetStore) ListUntestedAssetsWithSources(taskID int64, typ string, limit, offset int) ([]CoverageAsset, int, error) { + if limit <= 0 { + limit = 10 + } + if offset < 0 { + offset = 0 + } + typeFilter := "" + args := []any{taskID} + if typ != "" { + typeFilter = " AND target.type = $2" + args = append(args, typ) + } + var total int + if err := s.db.QueryRow(`WITH `+contextCoverageCTE+` +SELECT count(*) FROM target WHERE target.id NOT IN (SELECT asset_id FROM tested)`+typeFilter, args...).Scan(&total); err != nil { + return nil, 0, err + } + pageArgs := append(append([]any{}, args...), limit, offset) + limitPosition := strconv.Itoa(len(args) + 1) + offsetPosition := strconv.Itoa(len(args) + 2) + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` +SELECT target.id, target.type, target.label FROM target +WHERE target.id NOT IN (SELECT asset_id FROM tested)`+typeFilter+` +ORDER BY target.id LIMIT $`+limitPosition+` OFFSET $`+offsetPosition, pageArgs...) + if err != nil { + return nil, total, err + } + defer rows.Close() + out := []CoverageAsset{} + for rows.Next() { + var asset CoverageAsset + if err := rows.Scan(&asset.ID, &asset.Type, &asset.Label); err != nil { + return nil, total, err + } + out = append(out, asset) + } + return out, total, rows.Err() +} + +// HostsByTaskWithSources resolves exact HTTP host candidates from assets that +// are attached to, anchored by, or in scope for the current task or a direct +// source. Traffic remains global and is not copied. This read helper must not be +// used for destructive task cleanup; HostsByTask intentionally retains that +// narrower, task-owned behavior. +func (s *AssetStore) HostsByTaskWithSources(taskID int64) ([]string, error) { + rows, err := s.db.Query(`WITH `+directTaskContextCTE+`, +context_assets AS ( + SELECT DISTINCT a.id + FROM assets a + WHERE EXISTS (SELECT 1 FROM context_tasks ctx WHERE ctx.task_id=ANY(a.task_ids)) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes en ON en.id=ea.node_id + JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id + UNION + SELECT DISTINCT a.id + FROM assets a + JOIN task_scope ts ON ( + (ts.kind='company' AND a.company_id=ts.company_id) + OR (ts.kind='root_domain' AND a.root_domain=ts.domain) + OR (ts.kind='subdomain' AND a.domain=ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + ) + JOIN context_tasks ctx ON ctx.task_id=ts.task_id +) +SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,'') +FROM assets a JOIN context_assets ctx ON ctx.id=a.id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + hosts := map[string]struct{}{} + add := func(host string) { + host = strings.TrimSpace(strings.ToLower(host)) + if host != "" { + hosts[host] = struct{}{} + } + } + for rows.Next() { + var domain, ip, rawURL string + if err := rows.Scan(&domain, &ip, &rawURL); err != nil { + return nil, err + } + add(domain) + add(ip) + if parsed, err := url.Parse(rawURL); err == nil { + add(parsed.Hostname()) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]string, 0, len(hosts)) + for host := range hosts { + out = append(out, host) + } + sort.Strings(out) + return out, nil +} diff --git a/db/task_context.go b/db/task_context.go new file mode 100644 index 0000000..4927b02 --- /dev/null +++ b/db/task_context.go @@ -0,0 +1,388 @@ +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 + 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() +} + +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 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 WHERE id=$1`, taskID); err != nil { + return out, err + } + default: + return out, err + } + return out, tx.Commit() +} + +func truncateUTF8(value string, maxBytes int) string { + value = strings.ToValidUTF8(value, "\uFFFD") + if maxBytes <= 0 { + return "" + } + if len(value) <= maxBytes { + return value + } + end := maxBytes + for end > 0 && !utf8.ValidString(value[:end]) { + end-- + } + return value[:end] +} + +func (s *ExplorationStore) SetIntentBlockedReason(id int64, reason string) error { + _, err := s.db.Exec(`UPDATE exploration_nodes +SET state='blocked', blocked_reason=NULLIF($1,''), completed_at=now() +WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, reason, id, s.expID) + return err +} + +func (s *ExplorationStore) ReopenIntentsByBlockedReason(reason string) (int64, error) { + res, err := s.db.Exec(`UPDATE exploration_nodes +SET state='open', blocked_reason=NULL, completed_at=NULL +WHERE exploration_id=$1 AND kind='intent' AND state='blocked' AND blocked_reason=$2`, s.expID, reason) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + return n, nil +} diff --git a/db/task_context_lock_test.go b/db/task_context_lock_test.go new file mode 100644 index 0000000..06b455d --- /dev/null +++ b/db/task_context_lock_test.go @@ -0,0 +1,527 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" +) + +func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + first, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-first-%d", suffix), Format: "openai", Model: "first", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + second, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("lock-order-second-%d", suffix), Format: "openai", Model: "second", APIKey: "test-key", + }) + if err != nil { + _ = d.DeleteProfile(first) + t.Fatal(err) + } + task, err := d.CreateTaskWithOptions("LLM lock order", "verify concurrent mutation locks", TaskCreateOptions{ + LLMProfileIDs: []int64{first, second}, + }) + if err != nil { + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(first) + _ = d.DeleteProfile(second) + }) + + replaceDB, replacePID := openSingleConnectionTestDB(t, dsn) + replaceBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer replaceBlocker.Rollback() + if _, err := replaceBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + replaceDone := make(chan error, 1) + go func() { + replaceDone <- replaceDB.ReplaceTaskLLMProfiles(task.ID, []int64{second, first}, second) + }() + if err := waitForBackendBlock(d, replacePID, replaceDone); err != nil { + replaceBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, first, second) + if err := replaceBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(replaceDone); err != nil { + t.Fatalf("replace chain: %v", err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer deleteBlocker.Rollback() + if _, err := deleteBlocker.Exec(`SELECT id FROM tasks WHERE id=$1 FOR UPDATE`, task.ID); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + deleteDone := make(chan error, 1) + go func() { + deleteDone <- deleteDB.DeleteProfile(second) + }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + deleteBlocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, second) + if err := deleteBlocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != first { + t.Fatalf("deleting the active profile did not select its successor: %+v", got) + } +} + +func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-lock-profile-%d", suffix), Format: "openai", Model: "agent-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("lock_agent_%d", suffix), "lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + if err := d.SetAgentLLMProfile(agent.Key, &profileID); err != nil { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete agent profile: %v", err) + } + got, err := d.GetAgentByKey(agent.Key) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("agent binding was not cleared: agent=%+v err=%v", got, err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-lock-profile-%d", suffix), Format: "openai", Model: "conversation-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "lock conversation", &profileID) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete conversation profile: %v", err) + } + got, err := d.GetConversation(conversation.ID) + if err != nil || got == nil || got.LLMProfileID != nil { + t.Fatalf("conversation binding was not cleared: conversation=%+v err=%v", got, err) + } + }) +} + +func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) { + dsn := testDSN(t) + + t.Run("agent", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("agent-write-lock-profile-%d", suffix), Format: "openai", Model: "agent-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + agent, err := d.CreateAgent(fmt.Sprintf("write_lock_agent_%d", suffix), "write lock agent", "") + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteAgent(agent.Key) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.SetAgentLLMProfile(agent.Key, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind agent profile: %v", err) + } + }) + + t.Run("conversation", func(t *testing.T) { + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-write-lock-profile-%d", suffix), Format: "openai", Model: "conversation-write-lock", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + conversation, err := d.CreateConversation("planner", "write lock conversation", nil) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conversation.ID) + _ = d.DeleteProfile(profileID) + }) + + blocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer blocker.Rollback() + if _, err := blocker.Exec(`SELECT id FROM conversations WHERE id=$1 FOR UPDATE`, conversation.ID); err != nil { + t.Fatal(err) + } + mutationDB, mutationPID := openSingleConnectionTestDB(t, dsn) + mutationDone := make(chan error, 1) + go func() { mutationDone <- mutationDB.UpdateConversationProfile(conversation.ID, &profileID) }() + if err := waitForBackendBlock(d, mutationPID, mutationDone); err != nil { + blocker.Rollback() + t.Fatal(err) + } + assertProfilesUnlocked(t, d, profileID) + if err := blocker.Rollback(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(mutationDone); err != nil { + t.Fatalf("bind conversation profile: %v", err) + } + }) +} + +func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("conversation-create-race-profile-%d", suffix), Format: "openai", Model: "conversation-create-race", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + title := fmt.Sprintf("conversation create race %d", suffix) + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM conversations WHERE title=$1`, title) + _ = d.DeleteProfile(profileID) + }) + + profileBlocker, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer profileBlocker.Rollback() + if _, err := profileBlocker.Exec(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE`, profileID); err != nil { + t.Fatal(err) + } + + createDB, createPID := openSingleConnectionTestDB(t, dsn) + type createResult struct { + conversation *Conversation + err error + } + createDone := make(chan createResult, 1) + createStatus := make(chan error, 1) + go func() { + conversation, err := createDB.CreateConversation("planner", title, &profileID) + createDone <- createResult{conversation: conversation, err: err} + createStatus <- err + }() + if err := waitForBackendBlock(d, createPID, createStatus); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + profileBlocker.Rollback() + t.Fatal(err) + } + if err := profileBlocker.Rollback(); err != nil { + t.Fatal(err) + } + + var created createResult + select { + case created = <-createDone: + case <-time.After(12 * time.Second): + t.Fatal("timed out waiting for conversation creation") + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile during conversation creation: %v", err) + } + var persisted int + if err := d.QueryRow(`SELECT count(*) FROM conversations WHERE title=$1`, title).Scan(&persisted); err != nil { + t.Fatal(err) + } + if created.err == nil { + if created.conversation == nil || persisted != 1 { + t.Fatalf("successful creation was not committed atomically: conversation=%+v count=%d", created.conversation, persisted) + } + } else if persisted != 0 { + t.Fatalf("failed creation left a partial conversation row: err=%v count=%d", created.err, persisted) + } +} + +func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileID, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("late-reference-profile-%d", suffix), Format: "openai", Model: "late-reference", APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + task, err := d.CreateTask("late profile reference", "exercise delete retry", nil, 0, 0) + if err != nil { + _ = d.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _ = d.DeleteProfile(profileID) + }) + + // Keep the new reference uncommitted while deletion takes its initial + // READ COMMITTED snapshot. The task is therefore absent from the first lock + // set, while its FK KEY SHARE lock makes deletion wait at the profile row. + referenceTx, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer referenceTx.Rollback() + if _, err := referenceTx.Exec(`UPDATE tasks +SET llm_profile_id=$2, active_llm_profile_id=$2 +WHERE id=$1`, task.ID, profileID); err != nil { + t.Fatal(err) + } + if _, err := referenceTx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position) +VALUES ($1,$2,0)`, task.ID, profileID); err != nil { + t.Fatal(err) + } + + deleteDB, deletePID := openSingleConnectionTestDB(t, dsn) + deleteDone := make(chan error, 1) + go func() { deleteDone <- deleteDB.DeleteProfile(profileID) }() + if err := waitForBackendBlock(d, deletePID, deleteDone); err != nil { + referenceTx.Rollback() + t.Fatal(err) + } + if err := referenceTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForMutationResult(deleteDone); err != nil { + t.Fatalf("delete profile after late reference: %v", err) + } + + got, err := d.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if got.ActiveLLMProfileID != nil || len(got.LLMProfileIDs) != 0 || got.LLMChainRevision != 1 { + t.Fatalf("late task reference was not handled by a locked retry: %+v", got) + } +} + +func openSingleConnectionTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForBackendBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("mutation returned before reaching the expected reference-row lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func assertProfilesUnlocked(t *testing.T, d *DB, profileIDs ...int64) { + t.Helper() + probe, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer probe.Rollback() + for _, profileID := range profileIDs { + var lockedID int64 + if err := probe.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR UPDATE NOWAIT`, profileID).Scan(&lockedID); err != nil { + t.Fatalf("profile %d was locked before the task row: %v", profileID, err) + } + } +} + +func waitForMutationResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for task LLM mutation") + } +} diff --git a/db/task_context_unit_test.go b/db/task_context_unit_test.go new file mode 100644 index 0000000..519f304 --- /dev/null +++ b/db/task_context_unit_test.go @@ -0,0 +1,27 @@ +package db + +import ( + "strings" + "testing" + "unicode/utf8" +) + +func TestTruncateUTF8PreservesValidEncoding(t *testing.T) { + t.Parallel() + input := strings.Repeat("额度不足", 400) + got := truncateUTF8(input, 1000) + if len(got) > 1000 { + t.Fatalf("truncated value has %d bytes, want at most 1000", len(got)) + } + if !utf8.ValidString(got) { + t.Fatalf("truncated value is not valid UTF-8: %q", got[len(got)-8:]) + } +} + +func TestTruncateUTF8RepairsInvalidInput(t *testing.T) { + t.Parallel() + got := truncateUTF8("bad\xffvalue", 1000) + if !utf8.ValidString(got) { + t.Fatalf("repaired value is not valid UTF-8: %q", got) + } +} diff --git a/db/task_delete_concurrency_test.go b/db/task_delete_concurrency_test.go new file mode 100644 index 0000000..0d65191 --- /dev/null +++ b/db/task_delete_concurrency_test.go @@ -0,0 +1,210 @@ +package db + +import ( + "database/sql" + "fmt" + "testing" + "time" + + _ "github.com/jackc/pgx/v5/stdlib" +) + +func TestDeleteTaskTrafficHostsUseOneLockedTransaction(t *testing.T) { + dsn := testDSN(t) + d, err := Open(dsn) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + t.Run("sharing committed before deletion lock is observed", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/concurrent-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + writer, _ := openTaskDeleteTestDB(t, dsn) + writerTx, err := writer.Begin() + if err != nil { + t.Fatal(err) + } + defer writerTx.Rollback() + if _, err := writerTx.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID); err != nil { + t.Fatal(err) + } + + deleter, deleterPID := openTaskDeleteTestDB(t, dsn) + var preparedHosts []string + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + preparedHosts = append([]string(nil), p.TrafficHosts...) + return nil + }) + deleteDone <- deleteErr + }() + + if err := waitForTaskDeleteBlock(d, deleterPID, deleteDone); err != nil { + t.Fatal(err) + } + if err := writerTx.Commit(); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if containsDeleteHost(preparedHosts, host) { + t.Fatalf("newly shared host %q was selected for traffic deletion: %v", host, preparedHosts) + } + var remaining int + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE url=$1 AND $2=ANY(task_ids)`, serviceURL, second.ID).Scan(&remaining); err != nil { + t.Fatal(err) + } + if remaining != 1 { + t.Fatalf("concurrent owner's asset was not preserved: count=%d", remaining) + } + }) + + t.Run("sharing cannot enter between host resolution and commit", func(t *testing.T) { + first, second, host, rootAssetID := createTaskDeleteRaceFixture(t, d) + serviceURL := "https://" + host + "/late-owner" + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = d.Exec(`DELETE FROM assets WHERE id=$1 OR url=$2`, rootAssetID, serviceURL) + }) + + deleter, _ := openTaskDeleteTestDB(t, dsn) + prepared := make(chan TaskDeletePreparation, 1) + releasePrepare := make(chan struct{}) + deleteDone := make(chan error, 1) + go func() { + _, deleteErr := deleter.DeleteTaskCascadePrepared(first.ID, true, false, false, func(p TaskDeletePreparation) error { + prepared <- p + <-releasePrepare + return nil + }) + deleteDone <- deleteErr + }() + + var plan TaskDeletePreparation + select { + case plan = <-prepared: + case err := <-deleteDone: + close(releasePrepare) + t.Fatalf("delete returned before preparation: %v", err) + case <-time.After(5 * time.Second): + close(releasePrepare) + t.Fatal("timed out waiting for task delete preparation") + } + if !containsDeleteHost(plan.TrafficHosts, host) { + close(releasePrepare) + t.Fatalf("exclusive host %q missing from preparation: %v", host, plan.TrafficHosts) + } + + writer, writerPID := openTaskDeleteTestDB(t, dsn) + writerDone := make(chan error, 1) + go func() { + _, writeErr := writer.Exec(` +INSERT INTO assets(type, url, service_type, domain, task_ids) +VALUES ('service', $1, 'http', $2, ARRAY[$3]::bigint[])`, serviceURL, host, second.ID) + writerDone <- writeErr + }() + if err := waitForTaskDeleteBlock(d, writerPID, writerDone); err != nil { + close(releasePrepare) + t.Fatal(err) + } + + close(releasePrepare) + if err := waitForTaskDeleteResult(deleteDone); err != nil { + t.Fatal(err) + } + if err := waitForTaskDeleteResult(writerDone); err != nil { + t.Fatal(err) + } + }) +} + +func createTaskDeleteRaceFixture(t *testing.T, d *DB) (first, second *Task, host string, rootAssetID int64) { + t.Helper() + var err error + first, err = d.CreateTask("delete race owner", "delete safely", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err = d.CreateTask("delete race sharer", "preserve shared host", nil, 0, 0) + if err != nil { + _ = d.DeleteTask(first.ID) + t.Fatal(err) + } + host = fmt.Sprintf("task-delete-race-%d.example.test", first.ID) + rootAssetID, err = d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: host, TaskID: first.ID}) + if err != nil { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + t.Fatal(err) + } + return first, second, host, rootAssetID +} + +func openTaskDeleteTestDB(t *testing.T, dsn string) (*DB, int) { + t.Helper() + sqlDB, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDB.SetMaxIdleConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + if _, err := sqlDB.Exec(`SET statement_timeout='10s'`); err != nil { + t.Fatal(err) + } + var pid int + if err := sqlDB.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil { + t.Fatal(err) + } + return &DB{sqlDB}, pid +} + +func waitForTaskDeleteBlock(observer *DB, pid int, done <-chan error) error { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + select { + case err := <-done: + return fmt.Errorf("operation returned before reaching the task deletion lock: %v", err) + default: + } + var blockers int + if err := observer.QueryRow(`SELECT cardinality(pg_blocking_pids($1))`, pid).Scan(&blockers); err != nil { + return err + } + if blockers > 0 { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return fmt.Errorf("backend %d did not block within 5s", pid) +} + +func waitForTaskDeleteResult(done <-chan error) error { + select { + case err := <-done: + return err + case <-time.After(12 * time.Second): + return fmt.Errorf("timed out waiting for concurrent task deletion operation") + } +} + +func containsDeleteHost(hosts []string, want string) bool { + for _, host := range hosts { + if host == want { + return true + } + } + return false +} diff --git a/db/task_scope.go b/db/task_scope.go index e3108eb..2c69dcd 100644 --- a/db/task_scope.go +++ b/db/task_scope.go @@ -347,23 +347,23 @@ func hostPortOf(n *CoverageGraphNode) (string, int) { return host, port } -// BuildCoverageGraph assembles the full coverage graph for a task: every in-scope -// asset (all types) plus the connector root domains / companies needed to link -// them, with a tested flag on each in-scope asset. The frontend does the folding -// and highlighting; this only returns nodes + derived containment edges. -func (s *AssetStore) BuildCoverageGraph(taskID, expID int64) (*CoverageGraphData, error) { +// BuildCoverageGraph assembles the full coverage graph for a task and its direct +// read-only sources: every in-scope asset plus connector root domains/companies, +// with current-or-source fact anchors reflected in Tested. The legacy expID +// argument is retained for API compatibility; the task registry is authoritative. +func (s *AssetStore) BuildCoverageGraph(taskID, _ int64) (*CoverageGraphData, error) { g := &CoverageGraphData{Nodes: []CoverageGraphNode{}, Edges: []CoverageGraphEdge{}} if taskID <= 0 { return g, nil } - rows, err := s.db.Query(`WITH `+covTargetCTE+` + rows, err := s.db.Query(`WITH `+contextCoverageCTE+` SELECT a.id, a.type, COALESCE(a.company_id,0), COALESCE(a.domain,''), COALESCE(a.root_domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), COALESCE(a.port,0), COALESCE(a.service_type,''), COALESCE(a.app_name,''), COALESCE(a.page_title,''), COALESCE(a.status_code,0), (a.id IN (SELECT asset_id FROM tested)) AS tested FROM assets a JOIN target t ON t.id = a.id -ORDER BY a.id`, taskID, expID) +ORDER BY a.id`, taskID) if err != nil { return nil, err } diff --git a/db/task_source_limit_test.go b/db/task_source_limit_test.go new file mode 100644 index 0000000..01596c8 --- /dev/null +++ b/db/task_source_limit_test.go @@ -0,0 +1,20 @@ +package db + +import ( + "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) + } +} diff --git a/db/tasks.go b/db/tasks.go index 86dfac6..e5c4dc9 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -1,23 +1,34 @@ package db import ( + "database/sql" "encoding/json" + "fmt" + "strconv" "time" ) // Task is a row in the task registry (1:1 with an exploration). type Task struct { - ID int64 `json:"id"` - Description string `json:"description"` - Goal string `json:"goal"` - ExplorationID int64 `json:"exploration_id"` - Status string `json:"status"` - Paused bool `json:"paused"` - Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` - ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) - CreatedAt time.Time `json:"created_at"` - CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil + ID int64 `json:"id"` + Description string `json:"description"` + Goal string `json:"goal"` + ExplorationID int64 `json:"exploration_id"` + Status string `json:"status"` + Paused bool `json:"paused"` + Queued bool `json:"queued"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` + // Task-level ordered LLM chain. LLMProfileID remains the compatibility alias + // for ActiveLLMProfileID while older API clients still send one profile id. + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMChainRevision int64 `json:"-"` + LLMFailoverState string `json:"llm_failover_state,omitempty"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` + ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层) + CreatedAt time.Time `json:"created_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil // 任务级超时(见 docs/任务级超时与收尾设计.md)。 TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时 FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行 @@ -27,6 +38,23 @@ type Task struct { PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` } +// TaskDeleteResult reports optional related-data cleanup performed in the same +// transaction as the task/exploration delete. +type TaskDeleteResult struct { + AssetsDeleted int64 + AssetsDetached int64 + FindingsDeleted int64 + LLMRecordsDeleted int64 +} + +// TaskDeletePreparation is produced inside the PostgreSQL deletion transaction +// after asset and anchor writers have been excluded. Prepare callbacks may use +// TrafficHosts to stage an external traffic deletion before PostgreSQL commits. +type TaskDeletePreparation struct { + ExplorationID int64 + TrafficHosts []string +} + // IsTerminal reports whether a task status is a terminal (finished) state. // 单一真源，替换散落各处的 done/failed 硬编码判定。 func IsTerminal(status string) bool { @@ -40,6 +68,12 @@ func IsTerminal(status string) bool { // 低于它(含缺省 0 / 负值 / 误配的小值)一律抬到 10min，防止把 planner 打爆。 const MinPlanHeartbeatSeconds = 600 +// MaxTaskSourceCount bounds the amount of live inherited context one task can +// pull into every planner/main-agent prompt. Inheritance is intentionally direct +// only; keeping the fan-in bounded also prevents a single create request from +// multiplying graph and asset-context queries without limit. +const MaxTaskSourceCount = 8 + func normalizeHeartbeat(sec int) int { if sec < MinPlanHeartbeatSeconds { return MinPlanHeartbeatSeconds @@ -48,6 +82,30 @@ func normalizeHeartbeat(sec int) int { } func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { + var ids []int64 + if llmProfileID != nil { + ids = []int64{*llmProfileID} + } + return d.CreateTaskWithOptions(description, goal, TaskCreateOptions{ + LLMProfileIDs: ids, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds, + }) +} + +// TaskCreateOptions contains the task data that must be committed atomically +// with the task/exploration row. +type TaskCreateOptions struct { + SourceTaskIDs []int64 + LLMProfileIDs []int64 + TimeoutSeconds int + PlanHeartbeatSeconds int +} + +// CreateTaskWithOptions creates an exploration, task, direct source relations, +// and the ordered task LLM chain in one transaction. +func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOptions) (*Task, error) { + if len(opts.SourceTaskIDs) > MaxTaskSourceCount { + return nil, fmt.Errorf("too many source tasks: got %d, maximum is %d", len(opts.SourceTaskIDs), MaxTaskSourceCount) + } tx, err := d.Begin() if err != nil { return nil, err @@ -72,24 +130,80 @@ INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, or VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil { return nil, err } - if timeoutSeconds < 0 { - timeoutSeconds = 0 + if opts.TimeoutSeconds < 0 { + opts.TimeoutSeconds = 0 + } + opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds) + var active *int64 + if len(opts.LLMProfileIDs) > 0 { + id := opts.LLMProfileIDs[0] + active = &id + } + t := &Task{ + Description: description, Goal: goal, ExplorationID: expID, + LLMProfileID: active, ActiveLLMProfileID: active, + LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...), + SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...), + TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds, } - 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(` -INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6) -RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds) +VALUES ($1,$2,$3,$4,$4,$5,$6) +RETURNING id, status, paused, created_at`, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } + if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { + return nil, err + } + if err := 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() } -const taskCols = `id, description, goal, exploration_id, status, paused, queued, 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 insertTaskRelations(tx *sql.Tx, taskID int64, sourceIDs []int64) error { + seen := map[int64]bool{} + for _, sourceID := range sourceIDs { + if sourceID <= 0 || sourceID == taskID || seen[sourceID] { + return fmt.Errorf("invalid or duplicate source task id %d", sourceID) + } + seen[sourceID] = true + res, err := tx.Exec(`INSERT INTO task_relations(task_id, source_task_id) +SELECT $1, id FROM tasks WHERE id=$2 AND deleted_at IS NULL`, taskID, sourceID) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n != 1 { + return fmt.Errorf("source task %d not found", sourceID) + } + } + return nil +} + +func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { + seen := map[int64]bool{} + for position, profileID := range profileIDs { + if profileID <= 0 || seen[profileID] { + return fmt.Errorf("invalid or duplicate LLM profile id %d", profileID) + } + seen[profileID] = true + if _, err := tx.Exec(`INSERT INTO task_llm_profiles(task_id, profile_id, position) VALUES ($1,$2,$3)`, taskID, profileID, position); err != nil { + return err + } + } + return nil +} + +const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, active_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.Queued, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } return &t, nil @@ -116,7 +230,16 @@ func (d *DB) ListTasks() ([]*Task, error) { } out = append(out, t) } - return out, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + rows.Close() + for _, task := range out { + if err := d.hydrateTaskContext(task); err != nil { + return nil, err + } + } + return out, nil } // GetTask returns one alive task (nil if not found/deleted). @@ -128,6 +251,9 @@ func (d *DB) GetTask(id int64) (*Task, error) { } return nil, err } + if err := d.hydrateTaskContext(t); err != nil { + return nil, err + } return t, nil } @@ -137,8 +263,6 @@ func (d *DB) SetPaused(id int64, paused bool) error { return err } -// SetQueued persists a task's queued flag (并发上限挂起态). A queued task is held -// (engine not started) until the concurrency reconciler promotes it into a free slot. func (d *DB) SetQueued(id int64, queued bool) error { _, err := d.Exec(`UPDATE tasks SET queued=$1 WHERE id=$2`, queued, id) return err @@ -192,28 +316,132 @@ UPDATE tasks return deadline, err } -// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/ -// activity via ON DELETE CASCADE). The task row is removed first so the -// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete. -// Global assets are NOT touched. +// DeleteTask preserves the historical behavior: remove the task and its +// exploration subgraph while retaining global assets. func (d *DB) DeleteTask(id int64) error { + _, err := d.DeleteTaskCascade(id, false, false) + return err +} + +// DeleteTaskCascade hard-deletes a task and its exploration subgraph +// (nodes/edges/activity via ON DELETE CASCADE). Optional standalone findings and +// assets owned only by this task are deleted in the same transaction; shared +// assets are retained and only have this task id detached. deleteLLMRecords is +// optional to preserve callers of the original two-option API. +func (d *DB) DeleteTaskCascade(id int64, deleteAssets, deleteFindings bool, deleteLLMRecords ...bool) (TaskDeleteResult, error) { + deleteRecords := len(deleteLLMRecords) > 0 && deleteLLMRecords[0] + return d.DeleteTaskCascadePrepared(id, deleteAssets, deleteFindings, deleteRecords, nil) +} + +// DeleteTaskCascadePrepared coordinates reversible external deletion with the +// PostgreSQL cascade. When prepare is non-nil, host ownership is resolved and +// prepare is invoked inside this transaction while asset and anchor writers are +// excluded. The locks remain held through commit, closing the window where a +// host could become shared after its traffic had already been staged. +// +// prepare must only stage reversible work. Any returned error or later database +// error rolls PostgreSQL back; the caller remains responsible for rolling back +// external work that its callback staged successfully. +func (d *DB) DeleteTaskCascadePrepared( + id int64, + deleteAssets, deleteFindings, deleteLLMRecords bool, + prepare func(TaskDeletePreparation) error, +) (TaskDeleteResult, error) { + var result TaskDeleteResult tx, err := d.Begin() if err != nil { - return err + return result, err } defer tx.Rollback() var expID int64 - if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil { - if err.Error() == "sql: no rows in result set" { - return nil // already gone + if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 FOR UPDATE`, id).Scan(&expID); err != nil { + if err == sql.ErrNoRows { + return result, nil // already gone } - return err + return result, err + } + + // SHARE ROW EXCLUSIVE conflicts with every INSERT/UPDATE/DELETE on these + // tables and with another coordinated deletion. Taking both in one fixed order + // prevents phantoms (a distinct asset row for the same host) as well as new + // ownership/anchor references until the deletion transaction commits. + if deleteAssets || prepare != nil { + if _, err := tx.Exec(`LOCK TABLE assets, exploration_anchors IN SHARE ROW EXCLUSIVE MODE`); err != nil { + return result, err + } + } + if prepare != nil { + hosts, err := hostsForTaskDeletion(tx, id, expID) + if err != nil { + return result, err + } + if err := prepare(TaskDeletePreparation{ + ExplorationID: expID, + TrafficHosts: hosts, + }); err != nil { + return result, err + } + } + if deleteAssets { + // Ownership is the union of explicit task_ids and this exploration's anchors. + // An asset is deletable only when no other live task references it through + // either mechanism. This covers legacy anchor-only seeds without destroying + // evidence anchored by another task. + res, err := tx.Exec(` +WITH candidate_assets AS ( + SELECT id FROM assets WHERE $1 = ANY(task_ids) + UNION + SELECT ea.asset_id + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + WHERE n.exploration_id=$2 +), +deletable AS ( + SELECT a.id + FROM assets a + JOIN candidate_assets c ON c.id=a.id + WHERE NOT EXISTS ( + SELECT 1 FROM tasks t + WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids) + ) AND NOT EXISTS ( + SELECT 1 + FROM exploration_anchors ea + JOIN exploration_nodes n ON n.id=ea.node_id + JOIN tasks t ON t.exploration_id=n.exploration_id + WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL + ) +) +DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID) + if err != nil { + return result, err + } + result.AssetsDeleted, _ = res.RowsAffected() + + res, err = tx.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, id) + if err != nil { + return result, err + } + result.AssetsDetached, _ = res.RowsAffected() + } + if deleteFindings { + res, err := tx.Exec(`DELETE FROM findings WHERE task_id = $1`, id) + if err != nil { + return result, err + } + result.FindingsDeleted, _ = res.RowsAffected() + } + if deleteLLMRecords { + res, err := tx.Exec(`DELETE FROM llm_records WHERE COALESCE(task_id,'')=$1`, strconv.FormatInt(id, 10)) + if err != nil { + return result, err + } + result.LLMRecordsDeleted, _ = res.RowsAffected() } if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil { - return err + return result, err } if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil { - return err + return result, err } - return tx.Commit() + return result, tx.Commit() } diff --git a/db/tasks_test.go b/db/tasks_test.go index d407e03..655b06b 100644 --- a/db/tasks_test.go +++ b/db/tasks_test.go @@ -1,6 +1,11 @@ package db -import "testing" +import ( + "fmt" + "sync" + "testing" + "time" +) func TestTaskLifecycleAndDeleteCascade(t *testing.T) { d, err := Open(testDSN(t)) @@ -75,3 +80,406 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) { t.Fatalf("exploration should be deleted, got %d", exps) } } + +func TestTaskDeleteCascadeAssets(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + first, err := d.CreateTask("级联删除测试", "目标A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + second, err := d.CreateTask("共享资产保留测试", "目标B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + assets := d.Assets() + exclusiveID, err := assets.UpsertRootDomain(UpsertRootDomainReq{ + Domain: fmt.Sprintf("delete-%d.example.test", first.ID), + TaskID: first.ID, + }) + if err != nil { + t.Fatal(err) + } + sharedDomain := fmt.Sprintf("shared-%d.example.test", first.ID) + sharedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + if _, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: sharedDomain, TaskID: second.ID}); err != nil { + t.Fatal(err) + } + anchorOnlyDomain := fmt.Sprintf("anchor-only-%d.example.test", first.ID) + anchorOnlyID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: anchorOnlyDomain}) + if err != nil { + t.Fatal(err) + } + firstOrigin, err := d.Exploration(first.ExplorationID).OriginFactID() + if err != nil || firstOrigin == 0 { + t.Fatalf("first origin: id=%d err=%v", firstOrigin, err) + } + if err := d.Exploration(first.ExplorationID).Anchor(firstOrigin, anchorOnlyID); err != nil { + t.Fatal(err) + } + protectedDomain := fmt.Sprintf("other-anchor-%d.example.test", first.ID) + protectedID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: protectedDomain, TaskID: first.ID}) + if err != nil { + t.Fatal(err) + } + secondOrigin, err := d.Exploration(second.ExplorationID).OriginFactID() + if err != nil || secondOrigin == 0 { + t.Fatalf("second origin: id=%d err=%v", secondOrigin, err) + } + if err := d.Exploration(second.ExplorationID).Anchor(secondOrigin, protectedID); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(first.ID) + _ = d.DeleteTask(second.ID) + _, _ = assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + }) + + hosts, err := assets.HostsByTask(first.ID) + if err != nil { + t.Fatal(err) + } + hostSet := make(map[string]bool, len(hosts)) + for _, host := range hosts { + hostSet[host] = true + } + if !hostSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !hostSet[sharedDomain] { + t.Fatalf("task hosts missing cascade fixtures: %v", hosts) + } + deletableHosts, err := assets.HostsForTaskDeletion(first.ID, first.ExplorationID) + if err != nil { + t.Fatal(err) + } + deletableSet := make(map[string]bool, len(deletableHosts)) + for _, host := range deletableHosts { + deletableSet[host] = true + } + if !deletableSet[fmt.Sprintf("delete-%d.example.test", first.ID)] || !deletableSet[anchorOnlyDomain] { + t.Fatalf("exclusive or anchor-only cleanup host missing: %v", deletableHosts) + } + if deletableSet[sharedDomain] || deletableSet[protectedDomain] { + t.Fatalf("shared host was not protected from traffic cleanup: %v", deletableHosts) + } + + findingID, err := d.AddFinding(first.ID, 0, "__task_delete_cascade__", "", SeverityHigh, "summary", "evidence", "tester", []int64{exclusiveID}) + if err != nil { + t.Fatal(err) + } + + result, err := d.DeleteTaskCascade(first.ID, true, true) + if err != nil { + t.Fatal(err) + } + if result.AssetsDeleted < 2 || result.AssetsDetached < 2 { + t.Fatalf("unexpected asset cleanup result: %+v", result) + } + if result.FindingsDeleted != 1 { + t.Fatalf("unexpected finding cleanup result: %+v", result) + } + if finding, err := d.GetFinding(findingID); err != nil || finding != nil { + t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err) + } + remaining, err := assets.GetByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID}) + if err != nil { + t.Fatal(err) + } + if len(remaining) != 2 { + t.Fatalf("only shared and other-anchored assets should remain, got %+v", remaining) + } + byID := make(map[int64]*Asset, len(remaining)) + for _, asset := range remaining { + byID[asset.ID] = asset + } + if shared := byID[sharedID]; shared == nil || len(shared.TaskIDs) != 1 || shared.TaskIDs[0] != second.ID { + t.Fatalf("deleted task should be detached from shared asset: %+v", shared) + } + if protected := byID[protectedID]; protected == nil || len(protected.TaskIDs) != 0 { + t.Fatalf("other task's anchor should preserve the asset after detaching ownership: %+v", protected) + } +} + +func TestTaskRelationsAndLLMFailoverChain(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + suffix := time.Now().UnixNano() + profileIDs := make([]int64, 0, 3) + for i := 0; i < 3; i++ { + id, err := d.SaveProfile(&LLMProfile{ + Name: fmt.Sprintf("task-chain-%d-%d", suffix, i), Format: "openai", + Model: fmt.Sprintf("model-%d", i), APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + profileIDs = append(profileIDs, id) + } + + sourceA, err := d.CreateTask("source A", "goal A", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + sourceB, err := d.CreateTask("source B", "goal B", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + child, err := d.CreateTaskWithOptions("child", "new goal", TaskCreateOptions{ + SourceTaskIDs: []int64{sourceA.ID, sourceB.ID}, + LLMProfileIDs: profileIDs, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(child.ID) + _ = d.DeleteTask(sourceA.ID) + _ = d.DeleteTask(sourceB.ID) + for _, id := range profileIDs { + _ = d.DeleteProfile(id) + } + }) + + got, err := d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if fmt.Sprint(got.SourceTaskIDs) != fmt.Sprint([]int64{sourceA.ID, sourceB.ID}) { + t.Fatalf("unexpected direct sources: %v", got.SourceTaskIDs) + } + if fmt.Sprint(got.LLMProfileIDs) != fmt.Sprint(profileIDs) || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("unexpected initial chain: %+v", got) + } + initialRevision := got.LLMChainRevision + if err := d.ReplaceTaskLLMProfiles(child.ID, profileIDs, profileIDs[0]); err != nil { + t.Fatal(err) + } + staleRevision, err := d.MarkTaskLLMProfileQuotaExhaustedAtRevision(child.ID, profileIDs[0], initialRevision, "late quota from prior generation") + if err != nil { + t.Fatal(err) + } + if !staleRevision.Stale || staleRevision.Advanced || staleRevision.ChainExhausted { + t.Fatalf("old chain revision changed replacement chain: %+v", staleRevision) + } + got, err = d.GetTask(child.ID) + if err != nil { + t.Fatal(err) + } + if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] { + t.Fatalf("replacement revision/cursor not preserved: %+v", got) + } + + 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) + } + 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") + } +} diff --git a/llmrec/llmrec.go b/llmrec/llmrec.go index 37c256d..cf0be46 100644 --- a/llmrec/llmrec.go +++ b/llmrec/llmrec.go @@ -16,6 +16,27 @@ import ( "github.com/Autumn-27/artex/db" ) +type taskIDContextKey struct{} + +// WithTaskID attaches the owning task registry id to an LLM call. Session ids +// are based on exploration ids, which are not interchangeable with task ids. +func WithTaskID(ctx context.Context, taskID string) context.Context { + taskID = strings.TrimSpace(taskID) + if taskID == "" { + return ctx + } + return context.WithValue(ctx, taskIDContextKey{}, taskID) +} + +// TaskIDFrom returns the explicit task registry id attached by the task runtime. +func TaskIDFrom(ctx context.Context) string { + if ctx == nil { + return "" + } + taskID, _ := ctx.Value(taskIDContextKey{}).(string) + return strings.TrimSpace(taskID) +} + // Recorder wraps an llm.Provider and records every completion call. type Recorder struct { inner llm.Provider @@ -70,7 +91,13 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } start := time.Now() session := transcript.SessionIDFrom(ctx) - taskID, worker := parseSession(session) + parsedID, worker := parseSession(session) + taskID := TaskIDFrom(ctx) + if taskID == "" { + // Backward compatibility for non-task callers. For task calls, the task + // runtime always supplies the registry id explicitly. + taskID = parsedID + } // Serialize the request for storage. reqBody := serializeRequest(req) @@ -120,7 +147,9 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S } } -// record persists one LLM call to PostgreSQL (async, best-effort). +// record persists one LLM call to PostgreSQL before the provider stream returns. +// Keeping the write inside the owning task operation means task deletion can +// drain calls and then remove records without a late async insert recreating one. func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, reqBody string, start time.Time, text, thinking string, usage llm.Usage, stopReason string, streamErr error) { latency := int(time.Since(start).Milliseconds()) status := "ok" @@ -169,11 +198,9 @@ func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, re ResponseBody: string(respBody), } - go func() { - if err := r.pg.InsertLLMRecord(rec); err != nil { - log.Printf("[llmrec] insert: %v", err) - } - }() + if err := r.pg.InsertLLMRecord(rec); err != nil { + log.Printf("[llmrec] insert: %v", err) + } } // serializeRequest builds a JSON representation of the completion request. diff --git a/llmrec/llmrec_test.go b/llmrec/llmrec_test.go new file mode 100644 index 0000000..e998684 --- /dev/null +++ b/llmrec/llmrec_test.go @@ -0,0 +1,29 @@ +package llmrec + +import ( + "context" + "testing" +) + +func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) { + ctx := WithTaskID(context.Background(), " 42 ") + if got := TaskIDFrom(ctx); got != "42" { + t.Fatalf("TaskIDFrom()=%q want 42", got) + } + if got := TaskIDFrom(WithTaskID(ctx, " ")); got != "42" { + t.Fatalf("blank task id should preserve parent context, got %q", got) + } + if got := TaskIDFrom(nil); got != "" { + t.Fatalf("nil context returned %q", got) + } +} + +func TestParseSessionFallbackIsExplorationScoped(t *testing.T) { + taskID, worker := parseSession("exp12-worker-i99") + if taskID != "12" || worker != "worker" { + t.Fatalf("parseSession()=(%q,%q)", taskID, worker) + } + if taskID, worker := parseSession("not-a-task"); taskID != "" || worker != "" { + t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker) + } +} diff --git a/server/assets.go b/server/assets.go index 62a794e..ac8dcba 100644 --- a/server/assets.go +++ b/server/assets.go @@ -217,7 +217,6 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) { // GET /api/assets // ===================================================================== -// defaultAssetPageSize is the page size used when the caller omits ?limit. const defaultAssetPageSize = 50 func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { @@ -230,8 +229,6 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) { typ := q.Get("type") limit, _ := strconv.Atoi(q.Get("limit")) offset, _ := strconv.Atoi(q.Get("offset")) - // limit/offset are page controls, not a cap: `total` always reports the full - // match count so callers can page through everything. if limit <= 0 { limit = defaultAssetPageSize } @@ -351,8 +348,8 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) { Assets []struct { Type string `json:"type"` // root_domain / subdomain - Domain string `json:"domain"` - ICP string `json:"icp"` + Domain string `json:"domain"` + ICP string `json:"icp"` RecordType string `json:"record_type"` RecordValue []string `json:"record_value"` // ip diff --git a/server/chatupload.go b/server/chatupload.go index 7ac55e1..82ae080 100644 --- a/server/chatupload.go +++ b/server/chatupload.go @@ -45,9 +45,11 @@ type chatAttachment struct { // 文件先落这里,前端按返回的 abs 绝对路径写进任务描述) func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) { var sub string + taskScoped := false switch r.URL.Query().Get("scope") { case "task": sub = "tasks" + taskScoped = true case "session": sub = "sessions" case "staging": @@ -61,6 +63,17 @@ func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "非法 id") return } + if taskScoped { + if s.m.ResolveTask(id) == nil { + writeErr(w, 404, "task not found") + return + } + if !s.engine.beginTaskOperation(id) { + writeErr(w, http.StatusConflict, "任务正在删除，无法上传附件") + return + } + defer s.engine.decInflight(id) + } dir := filepath.Join(s.m.dir, sub, id, "uploads") if err := os.MkdirAll(dir, 0o755); err != nil { writeErr(w, 500, "建目录失败: "+err.Error()) diff --git a/server/core_test.go b/server/core_test.go index 49e9417..56b8a18 100644 --- a/server/core_test.go +++ b/server/core_test.go @@ -6,6 +6,8 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "os" + "path/filepath" "strconv" "testing" "time" @@ -92,15 +94,85 @@ func TestCoreTaskLifecyclePG(t *testing.T) { t.Fatalf("list tasks: %d", code) } - // delete → cascade removes the exploration subgraph - code, _ = doRetry("DELETE", "/api/tasks/"+id, nil) + // The task detail header consumes the top-level engine mode while some clients + // read the active-task snapshot. Keep both representations in sync. + code, out = doRetry("GET", "/api/stats?task="+id, nil) + if code != 200 { + t.Fatalf("task stats: %d (%v)", code, out) + } + activeTask, ok := out["active_task"].(map[string]any) + if !ok || activeTask["engine_mode"] != out["engine_mode"] { + t.Fatalf("engine mode mismatch: top=%v active_task=%v", out["engine_mode"], activeTask) + } + + // delete → cascade removes the exploration subgraph and selected related data + taskDir := filepath.Join(m.dir, "tasks", id) + if err := os.MkdirAll(taskDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(taskDir, "artifact.txt"), []byte("test"), 0o644); err != nil { + t.Fatal(err) + } + transcriptBase := "exp" + strconv.FormatInt(expID, 10) + "-main" + transcriptPath := filepath.Join(m.dir, "transcripts", transcriptBase+".jsonl") + sidechainPath := filepath.Join(m.dir, "transcripts", transcriptBase) + if err := os.MkdirAll(filepath.Join(sidechainPath, "subagents"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(transcriptPath, []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(sidechainPath, "subagents", "child.jsonl"), []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + nid, _ := strconv.ParseInt(id, 10, 64) + findingID, err := m.pg.AddFinding(nid, 0, "__core_task_delete__", "", db.SeverityHigh, "summary", "evidence", "tester", nil) + if err != nil { + t.Fatal(err) + } + if err := m.pg.EnsureLLMRecordsTable(); err != nil { + t.Fatal(err) + } + if err := m.pg.InsertLLMRecord(&db.LLMRecord{TaskID: id, SessionID: "delete-test", Status: "ok", RequestBody: "secret"}); err != nil { + t.Fatal(err) + } + // A non-canonical path still resolves to the one canonical task key for the + // barrier, workspace, in-memory registry and DB delete. + code, out = doRetry("DELETE", "/api/tasks/000"+id, map[string]bool{ + "delete_files": true, + "delete_findings": true, + "delete_llm_records": true, + }) if code != 200 { t.Fatalf("delete task: %d", code) } - nid, _ := strconv.ParseInt(id, 10, 64) + if deleted, _ := out["files_deleted"].(bool); !deleted { + t.Fatalf("expected files_deleted response, got %v", out) + } + if deleted, _ := out["findings_deleted"].(float64); deleted != 1 { + t.Fatalf("expected one deleted finding, got %v", out) + } + if deleted, _ := out["llm_records_deleted"].(float64); deleted != 1 { + t.Fatalf("expected one deleted LLM record, got %v", out) + } + if _, err := os.Stat(taskDir); !os.IsNotExist(err) { + t.Fatalf("task directory should be deleted, stat err=%v", err) + } + for _, path := range []string{transcriptPath, sidechainPath} { + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("task transcript should be deleted (%s), stat err=%v", path, err) + } + } if got, _ := m.pg.GetTask(nid); got != nil { t.Fatalf("task should be deleted from PG") } + if finding, err := m.pg.GetFinding(findingID); err != nil || finding != nil { + t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err) + } + var llmRecords int + if err := m.pg.QueryRow(`SELECT count(*) FROM llm_records WHERE task_id=$1`, id).Scan(&llmRecords); err != nil || llmRecords != 0 { + t.Fatalf("task LLM records should be deleted, count=%d err=%v", llmRecords, err) + } var n int m.pg.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, expID).Scan(&n) if n != 0 { diff --git a/server/dto.go b/server/dto.go index 1f7f0cd..9c320c0 100644 --- a/server/dto.go +++ b/server/dto.go @@ -30,21 +30,27 @@ func rawString(raw json.RawMessage) string { // ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status. type TaskDTO struct { - ID string `json:"id"` - ExplorationID int64 `json:"exploration_id"` - Description string `json:"description"` - Goal string `json:"goal"` - Status string `json:"status"` // created | running | paused | done | failed - CreatedAt string `json:"created_at"` - CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc) - CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished - CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished) - LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none) - Paused bool `json:"paused"` - Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption - GoalsTotal int `json:"goals_total"` - GoalsMet int `json:"goals_met"` - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default + ID string `json:"id"` + ExplorationID int64 `json:"exploration_id"` + Description string `json:"description"` + Goal string `json:"goal"` + Status string `json:"status"` // created | running | paused | done | failed + CreatedAt string `json:"created_at"` + CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc) + CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished + CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished) + LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none) + Paused bool `json:"paused"` + Queued bool `json:"queued"` + Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption + GoalsTotal int `json:"goals_total"` + GoalsMet int `json:"goals_met"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default + LLMProfileIDs []int64 `json:"llm_profile_ids"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMFailoverState string `json:"llm_failover_state"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []string `json:"source_task_ids"` } // TokenTotalDTO is a whole-task (all agents) token aggregate. @@ -65,18 +71,30 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO { } func taskDTO(t *Task, status string) TaskDTO { + llmState := t.llmStateSnapshot() + sourceIDs := make([]string, 0, len(t.SourceTaskIDs)) + for _, id := range t.SourceTaskIDs { + sourceIDs = append(sourceIDs, i64s(id)) + } + profileIDs := append(make([]int64, 0, len(llmState.ProfileIDs)), llmState.ProfileIDs...) return TaskDTO{ - ID: t.ID, - ExplorationID: t.ExpID, - Description: t.Description, - Goal: t.Goal, - Status: status, - CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)), - CreatedUnix: t.CreatedAt, - CompletedAt: completedRFC(t.CompletedAt), - CompletedUnix: t.CompletedAt, - Paused: t.Paused, - LLMProfileID: t.LLMProfileID, + ID: t.ID, + ExplorationID: t.ExpID, + Description: t.Description, + Goal: t.Goal, + Status: status, + CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)), + CreatedUnix: t.CreatedAt, + CompletedAt: completedRFC(t.CompletedAt), + CompletedUnix: t.CompletedAt, + Paused: t.Paused, + Queued: t.Queued, + LLMProfileID: llmState.ProfileID, + LLMProfileIDs: profileIDs, + ActiveLLMProfileID: llmState.ActiveID, + LLMFailoverState: llmState.FailoverState, + LLMFailoverReason: llmState.FailoverReason, + SourceTaskIDs: sourceIDs, } } @@ -120,17 +138,19 @@ func trafficDTOs(ex []traffic.ExchangeMeta) []TrafficExchangeDTO { // ---- TaskNode (frontend "TaskNode") — frontier, intents, exploration graph nodes ---- type TaskNodeDTO struct { - ID string `json:"id"` - Type string `json:"type"` // db Node.Kind - Payload string `json:"payload,omitempty"` - Priority int `json:"priority"` - State string `json:"state"` - Origin string `json:"origin"` - TS string `json:"ts"` + ID string `json:"id"` + Type string `json:"type"` // db Node.Kind + Payload string `json:"payload,omitempty"` + Priority int `json:"priority"` + State string `json:"state"` + Origin string `json:"origin"` + TS string `json:"ts"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` } func taskNodeDTO(n *db.Node) TaskNodeDTO { - return TaskNodeDTO{ + d := TaskNodeDTO{ ID: i64s(n.ID), Type: n.Kind, Payload: rawString(n.Payload), @@ -139,6 +159,11 @@ func taskNodeDTO(n *db.Node) TaskNodeDTO { Origin: n.Origin, TS: rfc3339(n.CreatedAt), } + if n.SourceTaskID > 0 { + d.SourceTaskID = i64s(n.SourceTaskID) + } + d.Inherited = n.Inherited + return d } func taskNodeDTOs(in []*db.Node) []TaskNodeDTO { @@ -169,6 +194,27 @@ func edgeDTOs(in []db.Edge) []EdgeDTO { return out } +// ---- Coverage asset references ---- + +type CoverageAssetRefDTO struct { + ID int64 `json:"id"` + Kind string `json:"kind"` + State string `json:"state"` + Summary string `json:"summary"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` +} + +func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO { + out := CoverageAssetRefDTO{ + ID: ref.ID, Kind: ref.Kind, State: ref.State, Summary: ref.Summary, Inherited: ref.Inherited, + } + if ref.SourceTaskID > 0 { + out.SourceTaskID = i64s(ref.SourceTaskID) + } + return out +} + // ---- Finding (frontend "Finding") ---- type FindingDTO struct { @@ -186,6 +232,8 @@ type FindingDTO struct { ParamID string `json:"param_id,omitempty"` TaskID string `json:"task_id,omitempty"` TaskDescription string `json:"task_description,omitempty"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` Assets []FindingAssetDTO `json:"assets,omitempty"` TS string `json:"ts"` } @@ -241,7 +289,7 @@ type findingPayload struct { func findingDTO(n *db.Node) FindingDTO { var p findingPayload _ = json.Unmarshal(n.Payload, &p) - return FindingDTO{ + d := FindingDTO{ ID: i64s(n.ID), VulnClass: p.VulnClass, Name: p.Name, @@ -251,6 +299,11 @@ func findingDTO(n *db.Node) FindingDTO { Evidence: rawString(p.Evidence), TS: rfc3339(n.CreatedAt), } + if n.SourceTaskID > 0 { + d.SourceTaskID = i64s(n.SourceTaskID) + } + d.Inherited = n.Inherited + return d } // findingDTOsForTask converts a task's finding nodes to DTOs, stamping each with @@ -260,11 +313,15 @@ func findingDTO(n *db.Node) FindingDTO { // nodes with no row keep the 'pending' default and no finding_id (not editable). // assets pre-resolves the anchored asset rows for label rendering. func findingDTOsForTask(t *Task, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO { + return findingDTOsForOwner(t.ID, t.Description, in, meta, assets) +} + +func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO { out := make([]FindingDTO, 0, len(in)) for _, n := range in { d := findingDTO(n) - d.TaskID = t.ID - d.TaskDescription = t.Description + d.TaskID = taskID + d.TaskDescription = description if m, ok := meta[n.ID]; ok { d.FindingID = i64s(m.ID) d.Status = m.Status @@ -320,16 +377,19 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO { // ---- Activity (frontend "Activity") ---- type ActivityDTO struct { - Seq int64 `json:"seq"` // db Activity.ID - IntentID string `json:"intent_id,omitempty"` // db NodeID - Worker string `json:"worker"` - TS string `json:"ts"` // db CreatedAt - Kind string `json:"kind"` - Tool string `json:"tool,omitempty"` - ToolUseID string `json:"tool_use_id,omitempty"` - IsError bool `json:"is_error"` - Summary string `json:"summary"` - Detail string `json:"detail,omitempty"` + Seq int64 `json:"seq"` // db Activity.ID + IntentID string `json:"intent_id,omitempty"` // db NodeID + Worker string `json:"worker"` + TS string `json:"ts"` // db CreatedAt + Kind string `json:"kind"` + Tool string `json:"tool,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + IsError bool `json:"is_error"` + Summary string `json:"summary"` + Detail string `json:"detail,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + SourceTaskID string `json:"source_task_id,omitempty"` + Inherited bool `json:"inherited,omitempty"` // token usage (set only on kind='result'); used for per-session token totals. InputTokens *int `json:"input_tokens,omitempty"` OutputTokens *int `json:"output_tokens,omitempty"` @@ -342,7 +402,7 @@ func activityDTO(a db.Activity) ActivityDTO { if a.NodeID != nil { intent = i64s(*a.NodeID) } - return ActivityDTO{ + d := ActivityDTO{ Seq: a.ID, IntentID: intent, Worker: a.Worker, @@ -353,11 +413,17 @@ func activityDTO(a db.Activity) ActivityDTO { IsError: a.IsError, Summary: a.Summary, Detail: a.Detail, // list endpoint leaves this empty (lazy) + Metadata: a.Metadata, InputTokens: a.InputTokens, OutputTokens: a.OutputTokens, CacheReadTokens: a.CacheReadTokens, CacheWriteTokens: a.CacheWriteTokens, } + if a.SourceTaskID > 0 { + d.SourceTaskID = i64s(a.SourceTaskID) + } + d.Inherited = a.Inherited + return d } func activityDTOs(in []db.Activity) []ActivityDTO { diff --git a/server/engine.go b/server/engine.go index 205f270..c52e76b 100644 --- a/server/engine.go +++ b/server/engine.go @@ -73,6 +73,13 @@ const ( modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避 ) +// retryableWorkerModelError excludes errors already handled by the task router. +// In particular, a quota error after partial streaming advances the task cursor +// for the next LLM call but must not replay this whole intent on the backup. +func retryableWorkerModelError(reason harness.TerminalReason, err error) bool { + return reason == harness.ReasonModelError && !isTaskLLMRuntimeError(err) +} + // Engine drives the event-driven exploration loop with real LLM agents // (docs §4.3/§4.4): on asset/exploration-graph change (debounced) it wakes the // planner, which reads the route, queries assets, judges goals and emits intents; @@ -86,28 +93,43 @@ type Engine struct { bc *Broadcaster // live activity pub/sub (SSE) - mu sync.RWMutex - planner *agent.Planner - worker *agent.Worker - started sync.Map // taskID -> bool, so Run is idempotent per task - lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat) - paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive) - dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records + mu sync.RWMutex + planner *agent.Planner + worker *agent.Worker + started sync.Map // taskID -> bool, so Run is idempotent per task + lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat) + llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls + paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive) + deleting sync.Map // taskID -> bool, delete barrier (no new task-owned writes) + // deleteWasPaused remembers whether BeginDelete temporarily paused an already + // running task. AbortDelete uses it to restore the exact pre-delete state. + deleteWasPaused sync.Map // taskID -> bool + dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records + + // deleteMu makes installing the delete barrier atomic with registering a new + // task operation. Once BeginDelete returns, every admitted writer is reflected + // in inflight and every later writer is rejected. + deleteMu sync.RWMutex + + // Every long-lived task goroutine (planner, workers and deadline coordinator) + // runs under one task-scoped context. Successful deletion cancels that context, + // waits for all goroutines, then releases every task-level Engine reference. + runtimeMu sync.Mutex + runtimes map[string]*taskRuntime // per-task execution context: each planner.Plan / worker.Execute runs under it, // so pausing can CANCEL an in-flight run (not just skip the next one). Recreated - // on resume since cancelling is one-shot. - // CancelCauseFunc (not CancelFunc): every cancel site attaches a named - // *agent.AbortCause so the run's trace can say WHICH of the ~10 cancel - // paths fired instead of one catch-all "运行被中断" line. + // on resume since cancelling is one-shot. Every cancellation carries a named + // cause so the activity trace can identify the initiating control path. execMu sync.Mutex execCancel map[string]context.CancelCauseFunc execCtx map[string]context.Context - // per-work (per running intent) cancel, so the planner's kill_work can stop a - // single worker without pausing the whole task. Keyed by globally-unique intent id. - workMu sync.Mutex - workCancel map[int64]context.CancelCauseFunc + // Per-work control lets the planner kill a worker and lets the UI pause/cancel + // one intent without pausing the whole task. The done channel closes only after + // runWorkerStep has stopped writing and committed its final state. + workMu sync.Mutex + work map[int64]*workExecution // steerBox queues planner course-corrections for a running work (keyed by intent // id). The worker's PreToolUse hook drains it before its next tool call and hands @@ -124,9 +146,23 @@ type Engine struct { inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain) coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重) - // resolve, if set, returns a task's dedicated planner/worker when it pins a - // specific LLM profile (nil,nil = use the global active pair). Set once at startup. - resolve func(t *Task) (*agent.Planner, *agent.Worker) + // resolve, if set, returns a task's dedicated planner/worker. In legacy mode + // nil,nil falls back to the global pair; an authoritative resolver uses nil,nil + // to mean this task is deliberately unavailable (for example an exhausted chain). + resolve func(t *Task) (*agent.Planner, *agent.Worker) + resolveAuthoritative bool +} + +type taskRuntime struct { + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup +} + +type workExecution struct { + cancel context.CancelCauseFunc + done chan struct{} + action string // user action: pause | cancel } // nextPlannerRound returns the next planner round number for a task (1-based). @@ -139,17 +175,132 @@ func (e *Engine) nextPlannerRound(taskID string) int { // Pause stops a task: marks it paused AND cancels any in-flight planner/worker run // for it (a long worker.Execute would otherwise keep going until it finishes). -// cause says WHY (user button / orchestrator tool / delete / reload) — it surfaces -// verbatim in the cancelled run's trace, so pass the matching agent.Abort* value. func (e *Engine) Pause(taskID string, cause error) { e.paused.Store(taskID, true) e.cancelExec(taskID, cause) } +// BeginDelete installs an execution barrier before task data/files are removed. +// The temporary pause is not a user pause: AbortDelete restores the exact state +// observed here when cleanup fails. +func (e *Engine) BeginDelete(taskID string) bool { + e.deleteMu.Lock() + if _, loaded := e.deleting.LoadOrStore(taskID, true); loaded { + e.deleteMu.Unlock() + return false + } + wasPaused := e.IsPaused(taskID) + e.deleteWasPaused.Store(taskID, wasPaused) + e.paused.Store(taskID, true) + e.deleteMu.Unlock() + e.cancelExec(taskID, agent.AbortTaskDeleted) + return true +} + +func (e *Engine) AbortDelete(taskID string) { + e.deleteMu.Lock() + if !e.IsDeleting(taskID) { + e.deleteMu.Unlock() + return + } + wasPaused := false + if v, ok := e.deleteWasPaused.Load(taskID); ok { + wasPaused, _ = v.(bool) + } + e.deleteWasPaused.Delete(taskID) + e.deleting.Delete(taskID) + if !wasPaused { + e.paused.Delete(taskID) + } + e.deleteMu.Unlock() + if !wasPaused && e.m != nil { + if t, ok := e.m.Task(taskID); ok { + t.Notify() + } + } +} + +func (e *Engine) IsDeleting(taskID string) bool { + _, ok := e.deleting.Load(taskID) + return ok +} + +// registerTaskRoutines reserves count goroutines in the task runtime. Callers +// hold deleteMu for reading so StopTask cannot race WaitGroup.Add with Wait. +func (e *Engine) registerTaskRoutines(parent context.Context, taskID string, count int) *taskRuntime { + e.runtimeMu.Lock() + defer e.runtimeMu.Unlock() + rt := e.runtimes[taskID] + if rt == nil { + ctx, cancel := context.WithCancel(parent) + rt = &taskRuntime{ctx: ctx, cancel: cancel} + e.runtimes[taskID] = rt + } + rt.wg.Add(count) + return rt +} + +func runTaskRoutine(rt *taskRuntime, fn func(context.Context)) { + go func() { + defer rt.wg.Done() + fn(rt.ctx) + }() +} + +// StopTask permanently stops every long-lived goroutine and removes all Engine +// state for a successfully deleted task. The delete barrier remains installed +// until cleanup finishes, so no new task operation can race the teardown. +func (e *Engine) StopTask(taskID string) { + e.deleteMu.Lock() + e.deleting.Store(taskID, true) + e.deleteMu.Unlock() + + e.cancelExec(taskID, agent.AbortTaskDeleted) + e.runtimeMu.Lock() + rt := e.runtimes[taskID] + if rt != nil { + rt.cancel() + } + e.runtimeMu.Unlock() + if rt != nil { + rt.wg.Wait() + } + + e.execMu.Lock() + if cancel := e.execCancel[taskID]; cancel != nil { + cancel(agent.AbortTaskDeleted) + } + delete(e.execCancel, taskID) + delete(e.execCtx, taskID) + e.execMu.Unlock() + + e.runtimeMu.Lock() + if e.runtimes[taskID] == rt { + delete(e.runtimes, taskID) + } + e.runtimeMu.Unlock() + + e.started.Delete(taskID) + e.lastAct.Delete(taskID) + e.llmCalls.Delete(taskID) + e.paused.Delete(taskID) + e.dropCnt.Delete(taskID) + e.plannerRound.Delete(taskID) + e.settling.Delete(taskID) + e.deadline.Delete(taskID) + e.stamped.Delete(taskID) + e.inflight.Delete(taskID) + e.coordStarted.Delete(taskID) + e.deleteWasPaused.Delete(taskID) + + e.deleteMu.Lock() + e.deleting.Delete(taskID) + e.deleteMu.Unlock() +} + // cancelExec cancels a task's current per-task exec context (any in-flight // planner.Plan / worker.Execute), if present. Shared by Pause and the settle -// sequence's hard-drain backstop. cause is recovered by the aborted run via -// context.Cause and printed in its result record. +// sequence's hard-drain backstop. func (e *Engine) cancelExec(taskID string, cause error) { e.execMu.Lock() if cancel := e.execCancel[taskID]; cancel != nil { @@ -206,32 +357,94 @@ func (e *Engine) LastActivity(taskID string) int64 { return 0 } +// BeginLLMCall/EndLLMCall track actual provider calls separately from the +// scheduler's task-operation counter. A task can have live loops while all of +// them are waiting for a trigger; that state must remain idle in the UI. +func (e *Engine) BeginLLMCall(taskID string) { + v, _ := e.llmCalls.LoadOrStore(taskID, new(int64)) + atomic.AddInt64(v.(*int64), 1) +} + +func (e *Engine) EndLLMCall(taskID string) { + if v, ok := e.llmCalls.Load(taskID); ok { + p := v.(*int64) + if atomic.AddInt64(p, -1) <= 0 { + atomic.StoreInt64(p, 0) + } + } +} + +func (e *Engine) ActiveLLMCalls(taskID string) int64 { + if v, ok := e.llmCalls.Load(taskID); ok { + return atomic.LoadInt64(v.(*int64)) + } + return 0 +} + func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) } func NewEngine(m *Manager) *Engine { return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(), execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{}, - workCancel: map[int64]context.CancelCauseFunc{}, steerBox: map[int64][]string{}} + work: map[int64]*workExecution{}, steerBox: map[int64][]string{}, + runtimes: map[string]*taskRuntime{}} } // registerWork records the cancel for the work currently running intentID. func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) { e.workMu.Lock() - e.workCancel[intentID] = cancel + e.work[intentID] = &workExecution{cancel: cancel, done: make(chan struct{})} e.workMu.Unlock() } -// unregisterWork removes and releases the work's cancel (called when Execute returns). -func (e *Engine) unregisterWork(intentID int64) { +// detachWork removes the live control handle once Execute has returned. complete +// must be called after the final intent state write so a waiting cancel handler can +// safely delete the worker's blackboard output without racing a late write. +func (e *Engine) detachWork(intentID int64) (action string, complete func()) { e.workMu.Lock() - if c := e.workCancel[intentID]; c != nil { - delete(e.workCancel, intentID) - c(agent.AbortWorkFinished) // release context resources (no-op if already cancelled) + run := e.work[intentID] + if run != nil { + delete(e.work, intentID) + action = run.action + run.cancel(agent.AbortWorkFinished) // release resources (no-op if already cancelled) } e.workMu.Unlock() e.steerMu.Lock() delete(e.steerBox, intentID) // drop any undelivered steering for a finished work e.steerMu.Unlock() + if run == nil { + return action, func() {} + } + return action, func() { close(run.done) } +} + +// ControlWork requests a user-visible pause or cancellation and waits until the +// worker has fully stopped writing. Cancellation cleanup is performed by the API +// handler after this returns; pause state is committed by runWorkerStep itself. +func (e *Engine) ControlWork(intentID int64, action string) error { + if action != "pause" && action != "cancel" { + return fmt.Errorf("unsupported work action %q", action) + } + e.workMu.Lock() + run := e.work[intentID] + if run == nil { + e.workMu.Unlock() + return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID) + } + if run.action != "" && run.action != action { + e.workMu.Unlock() + return fmt.Errorf("意图 %d 正在执行 %s 操作", intentID, run.action) + } + run.action = action + done := run.done + cause := error(agent.AbortWorkPausedByUser) + if action == "cancel" { + cause = agent.AbortWorkCancelledByUser + } + run.cancel(cause) + e.workMu.Unlock() + <-done + return nil } // SteerWork queues a mid-run course-correction for the work running intentID (the @@ -242,7 +455,7 @@ func (e *Engine) SteerWork(intentID int64, msg string) error { return fmt.Errorf("纠偏消息不能为空") } e.workMu.Lock() - running := e.workCancel[intentID] != nil + running := e.work[intentID] != nil e.workMu.Unlock() if !running { return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID) @@ -307,12 +520,12 @@ func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []s // The work's agent-core session honors ctx cancellation and aborts promptly. func (e *Engine) KillWork(intentID int64) error { e.workMu.Lock() - c := e.workCancel[intentID] + run := e.work[intentID] e.workMu.Unlock() - if c == nil { + if run == nil { return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID) } - c(agent.AbortKilledByPlanner) + run.cancel(agent.AbortKilledByPlanner) return nil } @@ -321,7 +534,7 @@ func (e *Engine) Broadcaster() *Broadcaster { return e.bc } // emitActivity persists one captured step AND fans it out to live subscribers, // from a single point so storage and the SSE stream never diverge. -func (e *Engine) emitActivity(t *Task, r db.Activity) { +func (e *Engine) emitActivity(t *Task, r db.Activity) db.Activity { id, err := e.appendActivity(t, r) if err != nil { // NO LONGER SILENT: dropping a record breaks command↔result pairing in the @@ -346,7 +559,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) { log.Printf("[activity] task %s 丢弃活动记录(该任务累计第 %d 条) worker=%s kind=%s tool=%s tuid=%s reason=%s summary=%q: %v%s", t.ID, n, r.Worker, r.Kind, r.Tool, r.ToolUseID, dropReason(err), preview(r.Summary, 80), err, diag) e.touch(t.ID) - return + return r } r.ID = id if r.CreatedAt.IsZero() { @@ -354,6 +567,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) { } e.bc.Publish(t.ID, r) e.touch(t.ID) + return r } // appendActivity persists one activity row, retrying briefly on write failure. @@ -396,6 +610,15 @@ func (e *Engine) snapshot() (*agent.Planner, *agent.Worker) { // Called once at startup before any task loop runs, so no lock is needed on reads. func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) { e.resolve = fn + e.resolveAuthoritative = false +} + +// SetAuthoritativeAgentResolver installs a resolver whose nil result must not +// fall through to the global provider. Task-level failover chains use this so a +// fully exhausted chain cannot silently bypass its configured boundary. +func (e *Engine) SetAuthoritativeAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) { + e.resolve = fn + e.resolveAuthoritative = true } // snapshotFor returns the planner/worker a task should run on: the task's dedicated @@ -403,7 +626,8 @@ func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worke // pair. The resolver returning nil means "no override" (fall back to global). func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) { if e.resolve != nil { - if p, w := e.resolve(t); p != nil && w != nil { + p, w := e.resolve(t) + if (p != nil && w != nil) || e.resolveAuthoritative { return p, w } } @@ -416,6 +640,14 @@ func (e *Engine) Ready() bool { return p != nil && w != nil } +// ReadyFor reports whether a specific task can resolve a planner/worker pair. +// An explicit task profile chain can be runnable even when no global default +// provider is configured, so task status must not rely on Ready alone. +func (e *Engine) ReadyFor(t *Task) bool { + p, w := e.snapshotFor(t) + return p != nil && w != nil +} + // Mode reports "llm" when ready, else "idle". func (e *Engine) Mode() string { if e.Ready() { @@ -428,15 +660,24 @@ func (e *Engine) Mode() string { // but no-op until an LLM is configured (so a task created while idle picks up // automatically once LLM is set from the UI). func (e *Engine) Run(ctx context.Context, t *Task) { + workers := e.m.Workers() + e.deleteMu.RLock() + if e.IsDeleting(t.ID) { + e.deleteMu.RUnlock() + return + } if _, loaded := e.started.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RUnlock() t.Notify() // already running — just nudge a planning round return } + rt := e.registerTaskRoutines(ctx, t.ID, 1+workers) + e.deleteMu.RUnlock() e.touch(t.ID) - go e.plannerLoop(ctx, t) - workers := e.m.Workers() + runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) }) for i := 0; i < workers; i++ { - go e.workerLoop(ctx, t, fmt.Sprintf("work#%d", i+1)) + name := fmt.Sprintf("work#%d", i+1) + runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, name) }) } e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重) // 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已 @@ -497,6 +738,9 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { if e.IsPaused(t.ID) { return // user-paused: don't plan } + if e.IsDeleting(t.ID) { + return + } // terminal task (goals all met → done, or failed): the run is over. A // resume/nudge — e.g. auto-resume of the active task on restart — must NOT // re-plan (it would burn an LLM round and re-confirm a settled result). @@ -508,10 +752,17 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { if e.isSettling(t.ID) { return } + if !e.beginTaskOperation(t.ID) { + return + } + defer e.decInflight(t.ID) e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) e.touch(t.ID) emit := func(r db.Activity) { e.emitActivity(t, r) } ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline + if ectx.Err() != nil || e.IsDeleting(t.ID) { + return + } log.Printf("[planner] task %s 规划中…(%s 触发)", t.ID, src) // round marker: each Plan() is one planner round; emit a boundary so the // UI can separate rounds in the transcript (kind='round'). @@ -520,10 +771,10 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { // what fired this round (worker done / finding; may be several — debounce // coalesces a burst; empty for time/heartbeat wakes). triggers := t.drainTriggers() - e.incInflight(t.ID) taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) met, reason, err := planner.Plan(ectx, taskIDInt, e.m.assets, t.Store, t.Goal, triggers, emit) - e.decInflight(t.ID) + e.EndLLMCall(t.ID) switch { case err != nil && ectx.Err() == nil: log.Printf("[planner] task %s 规划出错: %v", t.ID, err) @@ -578,6 +829,9 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { } continue // user-paused: don't claim/execute intents } + if e.IsDeleting(t.ID) { + return + } if e.isSettling(t.ID) { if sleepCtx(ctx, 1000*time.Millisecond) { return @@ -590,108 +844,144 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { } continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier } - intent := e.claimNext(t, name) - if intent == nil { - if sleepCtx(ctx, 800*time.Millisecond) { - return - } - continue - } - log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID) - e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) - e.touch(t.ID) - emit := func(r db.Activity) { e.emitActivity(t, r) } - ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline - // per-work child context so the planner's kill_work can stop just this work. - workCtx, workCancel := context.WithCancelCause(ectx) - e.registerWork(intent.ID, workCancel) - // wrap the guard hooks so steer_work can inject a mid-run course-correction - // for THIS intent (drained before the worker's next tool call). - iid := intent.ID - taskEmit := func(a db.Activity) { - nid := iid - a.NodeID, a.Worker = &nid, name - emit(a) - } - workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit) - hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }} - wTaskID, _ := strconv.ParseInt(t.ID, 10, 64) - e.incInflight(t.ID) // 计入在跑,供收尾时序 drain 等待 - reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) - // model_error 收场 → 额外重跑几次（退避后再试）。仅在意图仍属本 work、任务 - // 未暂停/未终止/未取消【且未进入收尾】时重试；否则让位给对应分支处理(收尾期不 - // 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。 - for attempt := 1; attempt <= modelErrorRetries && - reason == harness.ReasonModelError && - workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ { - log.Printf("[worker %s] task %s 意图 #%d model_error 收场，%v 后重试 (%d/%d)", - name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries) - if sleepCtx(workCtx, modelErrorRetryBackoff) { - break // 退避期间被取消（终止/暂停）→ 交给下方分支处理 - } - reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + if !e.beginTaskOperation(t.ID) { + return } + claimed := e.runWorkerStep(ctx, t, name, worker) e.decInflight(t.ID) - // CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's - // ctx was cancelled (planner kill_work) while the TASK ctx kept running; a - // pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER - // unregister would always be true (unregister cancels it) → every completed - // work would be wrongly marked stopped. - killed := workCtx.Err() != nil && ectx.Err() == nil - e.unregisterWork(intent.ID) - // if a pause cancelled this run mid-flight, return the intent to the frontier - // so it is re-claimed (and re-run from scratch) on resume — not marked done. - if ectx.Err() != nil && e.IsPaused(t.ID) { - _ = t.Store.SetIntentState(intent.ID, "open") - continue + if !claimed && sleepCtx(ctx, 800*time.Millisecond) { + return } - // 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾), - // 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。 - if ectx.Err() != nil && e.isSettling(t.ID) { - _ = t.Store.SetIntentState(intent.ID, "exhausted") - log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted)，写回 %s", name, t.ID, intent.ID, wrote) - e.touch(t.ID) - continue - } - // 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义, - // 标 stopped(不是 blocked),别污染已完成任务的意图状态。 - if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) { - _ = t.Store.SetIntentState(intent.ID, "stopped") - log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID) - e.touch(t.ID) - continue - } - // killed by the planner: mark stopped (don't write back results, don't auto-reclaim). - if killed { - _ = t.Store.SetIntentState(intent.ID, "stopped") - log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID) - e.touch(t.ID) - t.Notify() - continue - } - if err != nil { - log.Printf("[worker %s] intent %d: %v", name, intent.ID, err) - } - // terminal分流：撞步数上限 ≠ 完成。max_turns→exhausted（规划者据此知道这个方向 - // 试过但没真正做完、需换角度，而非当成已覆盖永久跳过）；出错→blocked；正常→done。 - state := "done" - switch { - case err != nil: - state = "blocked" - case reason == harness.ReasonMaxTurns: - state = "exhausted" - log.Printf("[worker %s] intent %d 撞步数上限(exhausted)，本次写回 %s", name, intent.ID, wrote) - case reason == harness.ReasonTimeout: - state = "exhausted" - log.Printf("[worker %s] intent %d 运行超时(exhausted)，收尾后写回 %s", name, intent.ID, wrote) - } - _ = t.Store.SetIntentState(intent.ID, state) - log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote) - e.touch(t.ID) - t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id) } } +// runWorkerStep claims and fully settles at most one intent. Its caller holds one +// task-operation admission for the whole claim/execute/state-write sequence so a +// delete cannot observe quiescence between the LLM return and the final DB writes. +func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker *agent.Worker) bool { + intent := e.claimNext(t, name) + if intent == nil { + return false + } + log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID) + e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务) + e.touch(t.ID) + emit := func(r db.Activity) { e.emitActivity(t, r) } + ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline + if ectx.Err() != nil || e.IsDeleting(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "open") + return true + } + // per-work child context so the planner's kill_work can stop just this work. + workCtx, workCancel := context.WithCancelCause(ectx) + e.registerWork(intent.ID, workCancel) + // wrap the guard hooks so steer_work can inject a mid-run course-correction + // for THIS intent (drained before the worker's next tool call). + iid := intent.ID + taskEmit := func(a db.Activity) { + nid := iid + a.NodeID, a.Worker = &nid, name + emit(a) + } + workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit) + hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }} + wTaskID, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) + reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + e.EndLLMCall(t.ID) + // model_error 收场 → 额外重跑几次（退避后再试）。仅在意图仍属本 work、任务 + // 未暂停/未终止/未取消【且未进入收尾】时重试；否则让位给对应分支处理(收尾期不 + // 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。 + for attempt := 1; attempt <= modelErrorRetries && + retryableWorkerModelError(reason, err) && + workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ { + log.Printf("[worker %s] task %s 意图 #%d model_error 收场，%v 后重试 (%d/%d)", + name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries) + if sleepCtx(workCtx, modelErrorRetryBackoff) { + break // 退避期间被取消（终止/暂停）→ 交给下方分支处理 + } + e.BeginLLMCall(t.ID) + reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding) + e.EndLLMCall(t.ID) + } + // Capture kill state before detachWork cancels workCtx. kill = this work's + // ctx was cancelled (planner kill_work) while the TASK ctx kept running; a + // pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER + // unregister would always be true (unregister cancels it) → every completed + // work would be wrongly marked stopped. + killed := workCtx.Err() != nil && ectx.Err() == nil + action, completeWork := e.detachWork(intent.ID) + defer completeWork() + if action == "pause" { + _ = t.Store.SetIntentState(intent.ID, "paused") + log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + if action == "cancel" { + // Leave the node in running until the waiting API handler deletes the intent + // and its direct outputs transactionally. No task Notify is emitted here. + log.Printf("[worker %s] task %s 意图 #%d 已停止，等待取消清理", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + // if a pause cancelled this run mid-flight, return the intent to the frontier + // so it is re-claimed (and re-run from scratch) on resume — not marked done. + if ectx.Err() != nil && e.IsPaused(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "open") + return true + } + // 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾), + // 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。 + if ectx.Err() != nil && e.isSettling(t.ID) { + _ = t.Store.SetIntentState(intent.ID, "exhausted") + log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted)，写回 %s", name, t.ID, intent.ID, wrote) + e.touch(t.ID) + return true + } + // 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义, + // 标 stopped(不是 blocked),别污染已完成任务的意图状态。 + if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) { + _ = t.Store.SetIntentState(intent.ID, "stopped") + log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID) + e.touch(t.ID) + return true + } + // killed by the planner: mark stopped (don't write back results, don't auto-reclaim). + if killed { + _ = t.Store.SetIntentState(intent.ID, "stopped") + log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID) + e.touch(t.ID) + t.Notify() + return true + } + if err != nil { + log.Printf("[worker %s] intent %d: %v", name, intent.ID, err) + } + // terminal分流：撞步数上限 ≠ 完成。max_turns→exhausted（规划者据此知道这个方向 + // 试过但没真正做完、需换角度，而非当成已覆盖永久跳过）；出错→blocked；正常→done。 + state := "done" + switch { + case err != nil: + state = "blocked" + case reason == harness.ReasonMaxTurns: + state = "exhausted" + log.Printf("[worker %s] intent %d 撞步数上限(exhausted)，本次写回 %s", name, intent.ID, wrote) + case reason == harness.ReasonTimeout: + state = "exhausted" + log.Printf("[worker %s] intent %d 运行超时(exhausted)，收尾后写回 %s", name, intent.ID, wrote) + } + if state == "blocked" && isTaskLLMChainExhausted(err) { + _ = t.Store.SetIntentBlockedReason(intent.ID, db.IntentBlockedLLMQuota) + } else { + _ = t.Store.SetIntentState(intent.ID, state) + } + log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote) + e.touch(t.ID) + t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id) + return true +} + func sleepCtx(ctx context.Context, d time.Duration) (done bool) { select { case <-ctx.Done(): diff --git a/server/engine_cancelcause_test.go b/server/engine_cancelcause_test.go index 0669679..5501821 100644 --- a/server/engine_cancelcause_test.go +++ b/server/engine_cancelcause_test.go @@ -7,90 +7,103 @@ import ( "github.com/Autumn-27/artex/agent" ) -// The engine's cancel paths must hand a NAMED cause to the run being cancelled — -// that cause is the whole content of the "运行被中断" line in the trace. These -// exercise the context machinery only (no Manager / DB needed). - func TestPauseCancelsWithNamedCause(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - // the per-work child the worker loop actually runs under - workCtx, workCancel := context.WithCancelCause(ectx) + execCtx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(execCtx) defer workCancel(nil) - e.Pause("t1", agent.AbortPausedByUser) - if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "paused_by_user" { - t.Fatalf("work ctx 拿到的原因 code=%q ok=%v，期望 paused_by_user", code, ok) + t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok) } } func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - workCtx, workCancel := context.WithCancelCause(ectx) + execCtx := e.execContextFor(context.Background(), "t1") + workCtx, workCancel := context.WithCancelCause(execCtx) e.registerWork(42, workCancel) - if err := e.KillWork(42); err != nil { - t.Fatalf("KillWork: %v", err) + t.Fatal(err) } if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "killed_by_planner" { - t.Fatalf("code=%q ok=%v，期望 killed_by_planner", code, ok) + t.Fatalf("code=%q ok=%v, want killed_by_planner", code, ok) } - // the task ctx must survive — kill_work stops one work, not the whole task - if ectx.Err() != nil { - t.Error("kill_work 不该取消整个任务的 exec ctx") - } - if err := e.KillWork(999); err == nil { - t.Error("对没有在跑的意图 KillWork 应该报错") + if execCtx.Err() != nil { + t.Fatal("kill_work cancelled the entire task context") } + _, complete := e.detachWork(42) + complete() } -func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) { +func TestControlWorkCarriesUserActionCause(t *testing.T) { for _, tc := range []struct { - name string - cause *agent.AbortCause - want string + action string + code string }{ - {"settle", agent.AbortSettleDrainTimeout, "settle_drain_timeout"}, - {"goalMet", agent.AbortGoalMet, "goal_met"}, + {action: "pause", code: "work_paused_by_user"}, + {action: "cancel", code: "work_cancelled_by_user"}, } { - t.Run(tc.name, func(t *testing.T) { + t.Run(tc.action, func(t *testing.T) { e := NewEngine(nil) - ectx := e.execContextFor(context.Background(), "t1") - e.cancelExec("t1", tc.cause) - if code, _, _, _ := agent.AbortReason(ectx); code != tc.want { - t.Fatalf("code=%q，期望 %q", code, tc.want) + ctx, cancel := context.WithCancelCause(context.Background()) + e.registerWork(42, cancel) + done := make(chan error, 1) + go func() { done <- e.ControlWork(42, tc.action) }() + <-ctx.Done() + if code, _, _, _ := agent.AbortReason(ctx); code != tc.code { + t.Fatalf("code=%q, want %q", code, tc.code) + } + action, complete := e.detachWork(42) + if action != tc.action { + t.Fatalf("action=%q, want %q", action, tc.action) + } + complete() + if err := <-done; err != nil { + t.Fatal(err) } }) } } -// asking for an exec ctx while paused must return an already-cancelled ctx that -// still SAYS why — this is the claim→Execute race guard, easily mistaken for a -// mystery abort. +func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) { + for _, tc := range []struct { + cause *agent.AbortCause + code string + }{ + {agent.AbortSettleDrainTimeout, "settle_drain_timeout"}, + {agent.AbortGoalMet, "goal_met"}, + } { + e := NewEngine(nil) + ctx := e.execContextFor(context.Background(), "t1") + e.cancelExec("t1", tc.cause) + if code, _, _, _ := agent.AbortReason(ctx); code != tc.code { + t.Fatalf("code=%q, want %q", code, tc.code) + } + } +} + func TestExecContextWhilePausedCarriesRaceGuardCause(t *testing.T) { e := NewEngine(nil) e.Pause("t1", agent.AbortPausedByUser) ctx := e.execContextFor(context.Background(), "t1") if ctx.Err() == nil { - t.Fatal("暂停期间不该发出可用的 exec ctx") + t.Fatal("paused task received a live execution context") } if code, _, _, _ := agent.AbortReason(ctx); code != "paused_race_guard" { - t.Fatalf("code=%q，期望 paused_race_guard", code) + t.Fatalf("code=%q, want paused_race_guard", code) } } -// Resume must hand out a live context again (cancelling is one-shot). func TestResumeYieldsFreshContext(t *testing.T) { e := NewEngine(nil) old := e.execContextFor(context.Background(), "t1") e.Pause("t1", agent.AbortPausedByUser) if old.Err() == nil { - t.Fatal("Pause 没取消旧 ctx") + t.Fatal("Pause did not cancel old context") } - e.paused.Delete("t1") // Resume() without the Task/Notify plumbing + e.paused.Delete("t1") if fresh := e.execContextFor(context.Background(), "t1"); fresh.Err() != nil { - t.Fatal("恢复后应拿到全新可用的 exec ctx") + t.Fatal("resume did not create a fresh execution context") } } diff --git a/server/engine_llm_calls_test.go b/server/engine_llm_calls_test.go new file mode 100644 index 0000000..56013e7 --- /dev/null +++ b/server/engine_llm_calls_test.go @@ -0,0 +1,43 @@ +package server + +import ( + "sync" + "testing" +) + +func TestEngineActiveLLMCallsConcurrent(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + const calls = 64 + + var started sync.WaitGroup + var release sync.WaitGroup + var finished sync.WaitGroup + started.Add(calls) + finished.Add(calls) + release.Add(1) + for range calls { + go func() { + defer finished.Done() + e.BeginLLMCall(taskID) + started.Done() + release.Wait() + e.EndLLMCall(taskID) + }() + } + started.Wait() + if got := e.ActiveLLMCalls(taskID); got != calls { + t.Fatalf("active calls = %d, want %d", got, calls) + } + release.Done() + finished.Wait() + if got := e.ActiveLLMCalls(taskID); got != 0 { + t.Fatalf("active calls after completion = %d, want 0", got) + } + + // A defensive extra end must never expose a negative UI count. + e.EndLLMCall(taskID) + if got := e.ActiveLLMCalls(taskID); got != 0 { + t.Fatalf("active calls after extra end = %d, want 0", got) + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 751c61d..7d98e27 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -43,7 +43,20 @@ func (e *Engine) inflightCounter(taskID string) *int64 { v, _ := e.inflight.LoadOrStore(taskID, new(int64)) return v.(*int64) } -func (e *Engine) incInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), 1) } + +// beginTaskOperation atomically registers a task-owned operation unless deletion +// has already installed its barrier. The delete handler can therefore wait for +// inflight==0 without a check-then-start race recreating files after cleanup. +func (e *Engine) beginTaskOperation(taskID string) bool { + e.deleteMu.RLock() + defer e.deleteMu.RUnlock() + if e.IsDeleting(taskID) { + return false + } + atomic.AddInt64(e.inflightCounter(taskID), 1) + return true +} + func (e *Engine) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) } func (e *Engine) inflightCount(taskID string) int64 { return atomic.LoadInt64(e.inflightCounter(taskID)) @@ -106,10 +119,18 @@ func (e *Engine) startDeadlineCoordinator(ctx context.Context, t *Task) { if t == nil || t.TimeoutSeconds <= 0 { return } - if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RLock() + if e.IsDeleting(t.ID) { + e.deleteMu.RUnlock() return } - go e.deadlineCoordinator(ctx, t) + if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded { + e.deleteMu.RUnlock() + return + } + rt := e.registerTaskRoutines(ctx, t.ID, 1) + e.deleteMu.RUnlock() + runTaskRoutine(rt, func(loopCtx context.Context) { e.deadlineCoordinator(loopCtx, t) }) } // deadlineCoordinator waits until the task's absolute deadline, then runs the settle @@ -170,6 +191,10 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。 met := e.runFinalPlannerRound(ctx, t) + if !e.beginTaskOperation(t.ID) { + return + } + defer e.decInflight(t.ID) // ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done, // 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。 @@ -192,6 +217,9 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // task-timeout planner words (final goal judgment; no new intents). Waits for the // LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout. func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { + if e.IsDeleting(t.ID) { + return false + } planner, _ := e.snapshotFor(t) for planner == nil { if sleepCtx(ctx, deadlinePollInterval) { @@ -200,17 +228,24 @@ func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { if isTerminalStatus(e.m.TaskStatus(t.ID)) { return false } + if e.IsDeleting(t.ID) { + return false + } planner, _ = e.snapshotFor(t) } + // 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。 + fctx := e.clockCtx(ctx, t, true) + if !e.beginTaskOperation(t.ID) { + return false + } + defer e.decInflight(t.ID) emit := func(r db.Activity) { e.emitActivity(t, r) } e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))}) - // 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。 - fctx := e.clockCtx(ctx, t, true) - e.incInflight(t.ID) tTaskID, _ := strconv.ParseInt(t.ID, 10, 64) + e.BeginLLMCall(t.ID) met, reason, err := planner.Plan(fctx, tTaskID, e.m.assets, t.Store, t.Goal, t.drainTriggers(), emit) - e.decInflight(t.ID) + e.EndLLMCall(t.ID) if err != nil { log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err) } else if met { diff --git a/server/goals.go b/server/goals.go index 41c4e99..3d83708 100644 --- a/server/goals.go +++ b/server/goals.go @@ -10,7 +10,6 @@ import ( "github.com/Autumn-27/artex/agent" "github.com/Autumn-27/artex/db" - "github.com/Autumn-27/norma/llm" ) type goalSpec struct { @@ -27,25 +26,36 @@ type goalSpec struct { // // 异步(goroutine)所以调用方立即返回,两条路径行为一致:秒建任务、后台拆目标。 func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) { + if !s.engine.beginTaskOperation(t.ID) { + return + } s.seed(t, seedText) if seedFirstIntent { s.seedFirstIntent(t) } - // 并发上限:满了就挂起排队(不做目标拆解、不启动引擎),由 reconcileConcurrency 补位启动。 if s.queueIfAtCapacity(t) { + s.engine.decInflight(t.ID) return } - go s.startTaskEngine(t) + go func() { + defer s.engine.decInflight(t.ID) + s.startTaskEngine(t) + }() } -// startTaskEngine does the actual running work: 第0轮目标拆解 + 启动引擎循环。 -// Shared by fresh launches and by the concurrency reconciler promoting a queued task. func (s *Server) startTaskEngine(t *Task) { + ctx := s.engine.execContextFor(s.ctx, t.ID) + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", Summary: "第 0 轮目标拆解"}) - goals := s.createGoals(s.ctx, t, func(r db.Activity) { + goals := s.createGoals(ctx, t, func(r db.Activity) { s.engine.emitActivity(t, r) }) + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } for _, g := range goals { summary := g.Text if g.VulnClass != "" { @@ -53,37 +63,26 @@ func (s *Server) startTaskEngine(t *Task) { } s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary}) } + if ctx.Err() != nil || s.engine.IsDeleting(t.ID) { + return + } s.engine.Run(s.ctx, t) } -// occupiesSlot reports whether a task counts against the concurrency cap: it is -// live (not terminal), not queued, and not paused. A just-created / decomposing -// task counts too (it will run momentarily). -func (s *Server) occupiesSlot(t *Task) bool { - if t == nil || t.Queued || isTerminalStatus(t.Status) { - return false - } - return !t.Paused && !s.engine.IsPaused(t.ID) +func (s *Server) occupiesConcurrencySlot(t *Task) bool { + return t != nil && !t.Queued && !t.Paused && !s.engine.IsPaused(t.ID) && !isTerminalStatus(t.Status) } -// runningTaskCount counts tasks occupying a concurrency slot, excluding excludeID -// (pass the candidate's id when deciding admission so it doesn't count itself). func (s *Server) runningTaskCount(excludeID string) int { - n := 0 - for _, t := range s.m.List() { - if t.ID == excludeID { - continue - } - if s.occupiesSlot(t) { - n++ + count := 0 + for _, task := range s.m.List() { + if task.ID != excludeID && s.occupiesConcurrencySlot(task) { + count++ } } - return n + return count } -// queueIfAtCapacity holds a task in the queued state when the concurrency cap is -// enabled and already full. Returns true if the task was queued (caller must not -// start it). No-op (false) when the feature is off or a slot is free. func (s *Server) queueIfAtCapacity(t *Task) bool { s.concMu.Lock() defer s.concMu.Unlock() @@ -93,52 +92,49 @@ func (s *Server) queueIfAtCapacity(t *Task) bool { } if err := s.m.SetTaskQueued(t.ID, true); err != nil { log.Printf("[concurrency] task %s 置排队失败: %v", t.ID, err) - return false // 落库失败宁可放它跑,也别静默丢任务 + return false } - t.Queued = true - s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", + s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: fmt.Sprintf("已排队：达到并发上限 %d，等待空位后自动开始", limit)}) - log.Printf("[concurrency] task %s 排队(并发上限 %d)", t.ID, limit) return true } -// reconcileConcurrency promotes queued tasks into free slots. Runs on every -// scheduler tick. When the feature is disabled it releases ALL queued tasks; -// when enabled it starts the oldest-queued tasks up to (limit - running). func (s *Server) reconcileConcurrency() { s.concMu.Lock() defer s.concMu.Unlock() - var queued []*Task - for _, t := range s.m.List() { - if t.Queued && !isTerminalStatus(t.Status) { - queued = append(queued, t) + queued := []*Task{} + for _, task := range s.m.List() { + if task.Queued && !isTerminalStatus(task.Status) { + queued = append(queued, task) } } - if len(queued) == 0 { - return - } - // 最早排队的先启动。 sort.Slice(queued, func(i, j int) bool { return queued[i].CreatedAt < queued[j].CreatedAt }) - enabled, limit := s.m.ConcurrencyLimit() - slots := len(queued) // feature off → 全部放行 + slots := len(queued) if enabled { - if !s.engine.Ready() { - return // 引擎未就绪(无 LLM):继续挂起,别空跑 - } slots = limit - s.runningTaskCount("") } - for _, t := range queued { + for _, task := range queued { if slots <= 0 { break } - if err := s.m.SetTaskQueued(t.ID, false); err != nil { - log.Printf("[concurrency] task %s 解除排队失败: %v", t.ID, err) + // Keep FIFO order, but do not consume a concurrency slot for a task + // whose explicit chain/global provider is unavailable. It will be retried + // after the user configures or resets its LLM chain. + if task.Paused || !s.engine.ReadyFor(task) { continue } - t.Queued = false - log.Printf("[concurrency] task %s 出队启动", t.ID) - go s.startTaskEngine(t) + if !s.engine.beginTaskOperation(task.ID) { + continue + } + if err := s.m.SetTaskQueued(task.ID, false); err != nil { + s.engine.decInflight(task.ID) + continue + } + go func(t *Task) { + defer s.engine.decInflight(t.ID) + s.startTaskEngine(t) + }(task) slots-- } } @@ -153,7 +149,6 @@ func (s *Server) reviveTask(t *Task) { if t == nil { return } - // 排队中的任务不复活启动——它本就等待并发空位,交给 reconcileConcurrency。 if t.Queued { return } @@ -184,24 +179,25 @@ func (s *Server) reviveTask(t *Task) { // Returns the seeded specs so callers can emit activity records for them. // emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI. func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec { + if t == nil { + return nil + } // Use the SAME LLM the task runs on (its pinned profile, else the active profile), // NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars, // so FromEnv returned empty and every task silently fell back to the crude rule // splitter (which shredded URLs / made meaningless 2-way splits). var specs []goalSpec - if prov, ok := s.goalsProvider(t); ok { - var as *db.AssetStore - if s.m != nil { - as = s.m.Assets() - } - taskID, _ := strconv.ParseInt(t.ID, 10, 64) - // DecomposeGoals persists each decomposed goal via the set_goals tool straight - // into t.Store; it returns the specs it wrote so we can emit activity + spot the - // empty (LLM error / no provider) case below. No write-back needed here anymore. - for _, g := range agent.DecomposeGoals(ctx, prov, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) { - if strings.TrimSpace(g.Text) != "" { - specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) - } + var as *db.AssetStore + if s.m != nil { + as = s.m.Assets() + } + taskID, _ := strconv.ParseInt(t.ID, 10, 64) + s.engine.BeginLLMCall(t.ID) + decomposed := agent.DecomposeGoalsWithProvider(ctx, s.agentsForTask(t).runtime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) + s.engine.EndLLMCall(t.ID) + for _, g := range decomposed { + if strings.TrimSpace(g.Text) != "" { + specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass}) } } if len(specs) == 0 { @@ -220,23 +216,3 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity } return specs } - -// goalsProvider resolves the provider for goal decomposition by the standard -// precedence: the goals agent's own binding → the task's pinned profile → the -// global active provider (the very instance the engine runs on, so decomposition -// shares its rate limiter, gets recorded, and follows LLM failover). -func (s *Server) goalsProvider(t *Task) (llm.Provider, bool) { - var pin *int64 - if t != nil { - pin = t.LLMProfileID - } - if eff := s.effectiveProfileForAgent("goals", pin); eff != nil { - if prov, cfg, ok := s.providerForProfile(*eff); ok { - return s.poolForBinding(*eff, prov, cfg), true - } - } - s.cfgMu.Lock() - prov, on := s.llmProv, s.llmOn - s.cfgMu.Unlock() - return prov, on && prov != nil -} diff --git a/server/goals_test.go b/server/goals_test.go new file mode 100644 index 0000000..057d85b --- /dev/null +++ b/server/goals_test.go @@ -0,0 +1,12 @@ +package server + +import ( + "context" + "testing" +) + +func TestCreateGoalsNilTask(t *testing.T) { + if got := (&Server{}).createGoals(context.Background(), nil, nil); len(got) != 0 { + t.Fatalf("nil task produced goals: %+v", got) + } +} diff --git a/server/inheritance_api_test.go b/server/inheritance_api_test.go new file mode 100644 index 0000000..df9fe4f --- /dev/null +++ b/server/inheritance_api_test.go @@ -0,0 +1,100 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Autumn-27/artex/db" +) + +func TestInheritedActivityDetailAndRelationDeletion(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + + source, err := m.CreateTaskWithOptions("detail source", "source goal", db.TaskCreateOptions{}) + if err != nil { + t.Fatal(err) + } + sourceID := mustTaskID(t, source.ID) + current, err := m.CreateTaskWithOptions("detail current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{sourceID}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(current.ID, DeleteTaskOptions{}) + _, _ = m.DeleteTask(source.ID, DeleteTaskOptions{}) + }) + + plannerStep, err := source.Store.AppendActivity(db.Activity{ + Worker: "planner", Kind: "text", Summary: "private planner", Detail: "private planner detail", + }) + if err != nil { + t.Fatal(err) + } + intentID, err := source.Store.AddIntent(map[string]any{"summary": "completed source work"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + if err := source.Store.SetIntentState(intentID, "done"); err != nil { + t.Fatal(err) + } + workerStep, err := source.Store.AppendActivity(db.Activity{ + NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "shared result", Detail: "shared worker detail", + }) + if err != nil { + t.Fatal(err) + } + + s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir()) + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatal(err) + } + fetchDetail := func(stepID int64) string { + req := httptest.NewRequest(http.MethodGet, + fmt.Sprintf("/api/exploration/activity/%d?task=%s", stepID, current.ID), nil) + req.Header.Set("Authorization", "Bearer "+token) + rec := httptest.NewRecorder() + s.Handler().ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("activity detail %d: status=%d body=%s", stepID, rec.Code, rec.Body.String()) + } + var body struct { + Detail string `json:"detail"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + return body.Detail + } + if got := fetchDetail(plannerStep); got != "" { + t.Fatalf("source planner transcript leaked through API: %q", got) + } + if got := fetchDetail(workerStep); got != "shared worker detail" { + t.Fatalf("terminal inherited worker detail missing: %q", got) + } + + if _, err := m.DeleteTask(source.ID, DeleteTaskOptions{}); err != nil { + t.Fatal(err) + } + currentAfter, ok := m.Task(current.ID) + if !ok || len(currentAfter.SourceTaskIDs) != 0 { + t.Fatalf("deleted source remained in live task DTO state: %+v", currentAfter) + } +} + +func mustTaskID(t *testing.T, raw string) int64 { + t.Helper() + var id int64 + if _, err := fmt.Sscan(raw, &id); err != nil || id <= 0 { + t.Fatalf("bad task id %q: %v", raw, err) + } + return id +} diff --git a/server/inheritance_dto_test.go b/server/inheritance_dto_test.go new file mode 100644 index 0000000..078d233 --- /dev/null +++ b/server/inheritance_dto_test.go @@ -0,0 +1,117 @@ +package server + +import ( + "encoding/json" + "testing" + "time" + + "github.com/Autumn-27/artex/db" +) + +func TestInheritedDTOProvenance(t *testing.T) { + createdAt := time.Date(2026, time.August, 19, 12, 0, 0, 0, time.UTC) + node := &db.Node{ + ID: 41, + Kind: db.KindFinding, + Payload: json.RawMessage(`{"vulnclass":"xss","severity":"high","summary":"source finding"}`), + State: "confirmed", + CreatedAt: createdAt, + SourceTaskID: 7, + Inherited: true, + } + + nodeDTO := taskNodeDTO(node) + if !nodeDTO.Inherited || nodeDTO.SourceTaskID != "7" { + t.Fatalf("node provenance lost: %+v", nodeDTO) + } + finding := findingDTOsForOwner("7", "source", []*db.Node{node}, nil, nil) + if len(finding) != 1 || !finding[0].Inherited || finding[0].SourceTaskID != "7" || finding[0].TaskID != "7" { + t.Fatalf("finding provenance lost: %+v", finding) + } + + intentID := int64(44) + activity := activityDTO(db.Activity{ + ID: 45, + NodeID: &intentID, + CreatedAt: createdAt, + SourceTaskID: 7, + Inherited: true, + }) + if !activity.Inherited || activity.SourceTaskID != "7" || activity.IntentID != "44" { + t.Fatalf("activity provenance lost: %+v", activity) + } + assetRef := coverageAssetRefDTO(db.AssetRef{ + ID: 46, Kind: db.KindFact, State: "confirmed", Summary: "source fact", SourceTaskID: 7, Inherited: true, + }) + if !assetRef.Inherited || assetRef.SourceTaskID != "7" { + t.Fatalf("asset ref provenance lost: %+v", assetRef) + } +} + +func TestInheritedIntentResultOnlyAllowsHistory(t *testing.T) { + for _, state := range []string{"done", "blocked", "exhausted", "stopped"} { + if !inheritedIntentResult(state) { + t.Errorf("terminal state %q should be inherited history", state) + } + } + for _, state := range []string{"open", "running", ""} { + if inheritedIntentResult(state) { + t.Errorf("live state %q must not enter inherited history", state) + } + } +} + +func TestInheritedGraphSnapshotFiltersLiveIntentsAndEdges(t *testing.T) { + nodes := []*db.Node{ + {ID: 1, Kind: db.KindFact, State: "confirmed"}, + {ID: 2, Kind: db.KindIntent, State: "running"}, + {ID: 3, Kind: db.KindIntent, State: "done"}, + {ID: 4, Kind: db.KindFinding, State: "confirmed"}, + } + edges := []db.Edge{ + {From: 1, Rel: db.RelDerivedFrom, To: 2}, + {From: 1, Rel: db.RelDerivedFrom, To: 3}, + {From: 3, Rel: db.RelYields, To: 4}, + } + + gotNodes, gotEdges := inheritedGraphSnapshot(nodes, edges, 77) + if len(gotNodes) != 3 { + t.Fatalf("expected fact, terminal intent and finding; got %+v", gotNodes) + } + for _, node := range gotNodes { + if node.ID == 2 { + t.Fatal("running inherited intent leaked into graph snapshot") + } + if !node.Inherited || node.SourceTaskID != 77 { + t.Fatalf("source provenance missing: %+v", node) + } + } + if len(gotEdges) != 2 { + t.Fatalf("edge referencing hidden live intent was not removed: %+v", gotEdges) + } + for _, edge := range gotEdges { + if edge.From == 2 || edge.To == 2 { + t.Fatalf("live intent id leaked through edge: %+v", edge) + } + } +} + +func TestFindingProvenanceInTask(t *testing.T) { + task := &Task{ID: "10", SourceTaskIDs: []int64{7, 8}} + + local := int64(10) + if source, inherited, allowed := findingProvenanceInTask(task, &local); !allowed || inherited || source != "" { + t.Fatalf("local finding: source=%q inherited=%v allowed=%v", source, inherited, allowed) + } + sourceID := int64(7) + if source, inherited, allowed := findingProvenanceInTask(task, &sourceID); !allowed || !inherited || source != "7" { + t.Fatalf("direct source finding: source=%q inherited=%v allowed=%v", source, inherited, allowed) + } + unrelated := int64(6) + if _, _, allowed := findingProvenanceInTask(task, &unrelated); allowed { + t.Fatal("unrelated finding must not be readable through task context") + } + if _, _, allowed := findingProvenanceInTask(task, nil); allowed { + t.Fatal("detached finding must not be readable through task context") + } +} diff --git a/server/manager.go b/server/manager.go index 78e403d..5f8dc36 100644 --- a/server/manager.go +++ b/server/manager.go @@ -2,6 +2,7 @@ package server import ( "encoding/json" + "errors" "fmt" "log" "os" @@ -12,38 +13,44 @@ import ( "sync" "time" - actool "github.com/Autumn-27/norma/tool" "github.com/Autumn-27/artex/agent" pgdb "github.com/Autumn-27/artex/db" "github.com/Autumn-27/artex/enrich" "github.com/Autumn-27/artex/guard" "github.com/Autumn-27/artex/intercept" "github.com/Autumn-27/artex/traffic" + actool "github.com/Autumn-27/norma/tool" ) // Task is one engagement: a description + goal + its own exploration store, // sharing the process-wide asset store. ID is the PG task id as a string; ExpID // is the exploration the task owns. type Task struct { - ID string `json:"id"` - ExpID int64 `json:"exploration_id"` - Description string `json:"description"` - Goal string `json:"goal"` - CreatedAt int64 `json:"created_at"` - CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成 - Paused bool `json:"paused"` - Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 - ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录) - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置 - Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态；空/其它则由运行态推导) + ID string `json:"id"` + ExpID int64 `json:"exploration_id"` + Description string `json:"description"` + Goal string `json:"goal"` + CreatedAt int64 `json:"created_at"` + CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成 + Paused bool `json:"paused"` + Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑 + ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录) + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置 + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"` + LLMFailoverState string `json:"llm_failover_state,omitempty"` + LLMFailoverReason string `json:"llm_failover_reason,omitempty"` + SourceTaskIDs []int64 `json:"source_task_ids,omitempty"` + Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态；空/其它则由运行态推导) // 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。 - TimeoutSeconds int `json:"timeout_seconds"` - PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒) - FirstRunAt int64 `json:"first_run_at,omitempty"` - DeadlineAt int64 `json:"deadline_at,omitempty"` - Store *pgdb.ExplorationStore `json:"-"` - Guard *guard.Guard `json:"-"` - notify chan struct{} + TimeoutSeconds int `json:"timeout_seconds"` + PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒) + FirstRunAt int64 `json:"first_run_at,omitempty"` + DeadlineAt int64 `json:"deadline_at,omitempty"` + Store *pgdb.ExplorationStore `json:"-"` + Guard *guard.Guard `json:"-"` + notify chan struct{} + llmMu sync.RWMutex // pendingTriggers accumulates the concrete changes (worker done / finding) that // fired planning rounds since the last one consumed them. The debounce coalesces @@ -52,14 +59,74 @@ type Task struct { pendingTriggers []agent.TriggerEvent } +type taskLLMState struct { + ProfileID *int64 + ProfileIDs []int64 + ActiveID *int64 + FailoverState string + FailoverReason string +} + +func cloneInt64Ptr(value *int64) *int64 { + if value == nil { + return nil + } + copy := *value + return © +} + +func (t *Task) llmStateSnapshot() taskLLMState { + t.llmMu.RLock() + defer t.llmMu.RUnlock() + return taskLLMState{ + ProfileID: cloneInt64Ptr(t.LLMProfileID), + ProfileIDs: append(make([]int64, 0, len(t.LLMProfileIDs)), t.LLMProfileIDs...), + ActiveID: cloneInt64Ptr(t.ActiveLLMProfileID), + FailoverState: t.LLMFailoverState, + FailoverReason: t.LLMFailoverReason, + } +} + +func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, state, reason string) { + t.llmMu.Lock() + t.LLMProfileID = cloneInt64Ptr(profileID) + t.ActiveLLMProfileID = cloneInt64Ptr(activeID) + t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...) + t.LLMFailoverState = state + t.LLMFailoverReason = reason + t.llmMu.Unlock() +} + +// DeleteTaskOptions controls cleanup of data stored outside the task's own +// exploration graph. All options default to false for backward compatibility. +type DeleteTaskOptions struct { + DeleteAssets bool `json:"delete_assets"` + DeleteTraffic bool `json:"delete_traffic"` + DeleteFiles bool `json:"delete_files"` + DeleteFindings bool `json:"delete_findings"` + DeleteLLMRecords bool `json:"delete_llm_records"` +} + +// DeleteTaskResult makes destructive cleanup auditable to API callers. +type DeleteTaskResult struct { + Deleted string `json:"deleted"` + AssetsDeleted int64 `json:"assets_deleted"` + AssetsDetached int64 `json:"assets_detached"` + TrafficDeleted int64 `json:"traffic_deleted"` + FilesDeleted bool `json:"files_deleted"` + FindingsDeleted int64 `json:"findings_deleted"` + LLMRecordsDeleted int64 `json:"llm_records_deleted"` + CleanupWarning string `json:"cleanup_warning,omitempty"` +} + // Manager owns the PostgreSQL data source (asset graph + every task's exploration // graph + config) and the in-memory set of task handles. type Manager struct { dir string pg *pgdb.DB assets *pgdb.AssetStore - traffic *traffic.Traffic // process-wide recording proxy (may be nil) - enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP) + traffic *traffic.Traffic // process-wide recording proxy (may be nil) + enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP) interceptor *intercept.Interceptor // user-configured tool-call interception rules mu sync.RWMutex @@ -529,7 +596,10 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID, Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued, CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef, - LLMProfileID: pt.LLMProfileID, + LLMProfileID: pt.LLMProfileID, + LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID, + LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason, + SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...), TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1), @@ -539,7 +609,17 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter // CreateTask creates a task + its exploration and makes it active. // timeoutSeconds is the task-level wall-clock budget (0 = 不限时). func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) { - pt, err := m.pg.CreateTask(description, goal, llmProfileID, timeoutSeconds, planHeartbeatSeconds) + var ids []int64 + if llmProfileID != nil { + ids = []int64{*llmProfileID} + } + return m.CreateTaskWithOptions(description, goal, pgdb.TaskCreateOptions{ + LLMProfileIDs: ids, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds, + }) +} + +func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.TaskCreateOptions) (*Task, error) { + pt, err := m.pg.CreateTaskWithOptions(description, goal, opts) if err != nil { return nil, err } @@ -551,6 +631,35 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time return t, nil } +// ReplaceTaskLLMProfiles resets a non-terminal task's ordered provider chain and +// mirrors the committed state onto the live task handle. +func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activeProfileID int64) (int64, error) { + n, err := strconv.ParseInt(id, 10, 64) + if err != nil { + return 0, err + } + if err := m.pg.ReplaceTaskLLMProfiles(n, profileIDs, activeProfileID); err != nil { + return 0, err + } + pt, err := m.pg.GetTask(n) + if err != nil || pt == nil { + return 0, err + } + m.mu.Lock() + if task := m.tasks[id]; task != nil { + task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason) + } + m.mu.Unlock() + if task, ok := m.Task(id); ok { + reopened, reopenErr := task.Store.ReopenIntentsByBlockedReason(pgdb.IntentBlockedLLMQuota) + if reopenErr != nil { + return reopened, reopenErr + } + return reopened, nil + } + return 0, nil +} + // LoadExisting rebuilds in-memory task handles from the PG task registry. func (m *Manager) LoadExisting() []*Task { pts, err := m.pg.ListTasks() @@ -706,17 +815,123 @@ func (m *Manager) SetTaskStatus(id, status string) error { return nil } -// DeleteTask removes a task (and its exploration subgraph) from PG + memory. -func (m *Manager) DeleteTask(id string) error { +// DeleteTask removes a task and optionally its related global data. Traffic has +// no task-id column, so related exchanges are resolved by exact hosts from the +// task's asset rows. Files are staged before the database operation; traffic is +// staged while PostgreSQL excludes asset/anchor writers. Both are restored on a +// database failure and purged only after its commit. +func (m *Manager) DeleteTask(id string, opts DeleteTaskOptions) (DeleteTaskResult, error) { + result := DeleteTaskResult{Deleted: id} n, err := strconv.ParseInt(id, 10, 64) if err != nil { - return err + return result, err } - if err := m.pg.DeleteTask(n); err != nil { - return err + registered, err := m.pg.GetTask(n) + if err != nil { + return result, err } + if registered == nil { + m.forgetTask(id, n) + return result, nil + } + + var fileStage *taskFileDeleteStage + if opts.DeleteFiles { + fileStage, err = stageTaskFiles(m.dir, id, registered.ExplorationID) + if err != nil { + return result, err + } + result.FilesDeleted = fileStage.deleted + } + + var trafficStage *traffic.HostDeleteStage + var prepare func(pgdb.TaskDeletePreparation) error + if opts.DeleteTraffic && m.traffic != nil { + prepare = func(p pgdb.TaskDeletePreparation) error { + if len(p.TrafficHosts) == 0 { + return nil + } + trafficStage, err = m.traffic.StageDeleteHostsExact(p.TrafficHosts) + if err != nil { + return err + } + result.TrafficDeleted = trafficStage.Deleted() + return nil + } + } + + dbResult, err := m.pg.DeleteTaskCascadePrepared( + n, opts.DeleteAssets, opts.DeleteFindings, opts.DeleteLLMRecords, prepare, + ) + if err != nil { + return result, rollbackTaskDelete(err, trafficStage, fileStage) + } + result.AssetsDeleted = dbResult.AssetsDeleted + result.AssetsDetached = dbResult.AssetsDetached + result.FindingsDeleted = dbResult.FindingsDeleted + result.LLMRecordsDeleted = dbResult.LLMRecordsDeleted + + // PostgreSQL is now authoritative: finalize the staged external deletion and + // forget the live task even if a final purge reports an error. Such errors are + // typed so the HTTP layer can still tear down the task runtime instead of + // incorrectly reviving a task whose database row is already gone. + var finalizeErrs []error + if trafficStage != nil { + if err := trafficStage.Commit(); err != nil { + finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize traffic deletion: %w", err)) + } + } + if fileStage != nil { + if err := fileStage.commit(); err != nil { + finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize task file deletion: %w", err)) + } + } + m.forgetTask(id, n) + if err := errors.Join(finalizeErrs...); err != nil { + return result, &taskDeleteCommittedError{err: err} + } + return result, nil +} + +// taskDeleteCommittedError means PostgreSQL deletion succeeded but purging one +// of the recoverable staging directories failed. The task must stay deleted. +type taskDeleteCommittedError struct{ err error } + +func (e *taskDeleteCommittedError) Error() string { + return "task deletion committed; external cleanup incomplete: " + e.err.Error() +} + +func (e *taskDeleteCommittedError) Unwrap() error { return e.err } + +func rollbackTaskDelete(cause error, trafficStage *traffic.HostDeleteStage, fileStage *taskFileDeleteStage) error { + errs := []error{cause} + // Reverse the preparation order. Both restorations are attempted even if the + // first one fails, and errors.Join preserves the original PostgreSQL error. + if trafficStage != nil { + if err := trafficStage.Rollback(); err != nil { + errs = append(errs, fmt.Errorf("restore traffic after task delete failure: %w", err)) + } + } + if fileStage != nil { + if err := fileStage.rollback(); err != nil { + errs = append(errs, fmt.Errorf("restore task files after task delete failure: %w", err)) + } + } + return errors.Join(errs...) +} + +func (m *Manager) forgetTask(id string, numericID int64) { m.mu.Lock() delete(m.tasks, id) + for _, task := range m.tasks { + kept := task.SourceTaskIDs[:0] + for _, sourceID := range task.SourceTaskIDs { + if sourceID != numericID { + kept = append(kept, sourceID) + } + } + task.SourceTaskIDs = kept + } if m.active == id { m.active = "" for _, t := range m.tasks { @@ -725,7 +940,128 @@ func (m *Manager) DeleteTask(id string) error { } } m.mu.Unlock() - return nil +} + +type stagedTaskPath struct { + source string + staged string +} + +type taskFileDeleteStage struct { + stageDir string + moves []stagedTaskPath + deleted bool + done bool +} + +// stageTaskFiles atomically renames the task workspace and owned transcripts to +// a same-filesystem staging directory. The trailing dash in the transcript +// prefix is significant: exploration 12 must not match exploration 123. +func stageTaskFiles(dataDir, taskID string, explorationID int64) (*taskFileDeleteStage, error) { + stage := &taskFileDeleteStage{} + var targets []string + taskDir := filepath.Join(dataDir, "tasks", taskID) + if _, err := os.Lstat(taskDir); err == nil { + targets = append(targets, taskDir) + } else if !os.IsNotExist(err) { + return nil, err + } + + transcriptDir := filepath.Join(dataDir, "transcripts") + entries, err := os.ReadDir(transcriptDir) + if err != nil { + if !os.IsNotExist(err) { + return nil, err + } + entries = nil + } + prefix := fmt.Sprintf("exp%d-", explorationID) + for _, entry := range entries { + name := entry.Name() + if !strings.HasPrefix(name, prefix) || (!entry.IsDir() && !strings.HasSuffix(name, ".jsonl")) { + continue + } + targets = append(targets, filepath.Join(transcriptDir, name)) + } + if len(targets) == 0 { + stage.done = true + return stage, nil + } + + parent := filepath.Join(dataDir, ".delete-staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return nil, err + } + stage.stageDir, err = os.MkdirTemp(parent, "task-"+taskID+"-") + if err != nil { + return nil, err + } + for _, source := range targets { + staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source))) + if err := os.Rename(source, staged); err != nil { + cause := fmt.Errorf("stage task file %s: %w", source, err) + if restoreErr := stage.rollback(); restoreErr != nil { + return nil, errors.Join(cause, fmt.Errorf("restore partially staged task files: %w", restoreErr)) + } + return nil, cause + } + stage.moves = append(stage.moves, stagedTaskPath{source: source, staged: staged}) + } + stage.deleted = true + return stage, nil +} + +func (s *taskFileDeleteStage) commit() error { + if s == nil || s.done { + return nil + } + err := os.RemoveAll(s.stageDir) + s.done = true + return err +} + +func (s *taskFileDeleteStage) rollback() error { + if s == nil || s.done { + return nil + } + var errs []error + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.MkdirAll(filepath.Dir(move.source), 0o755); err != nil { + errs = append(errs, fmt.Errorf("create restore parent for %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove task file stage: %w", err)) + } + } + s.done = true + return errors.Join(errs...) +} + +// deleteTaskFiles retains the standalone helper contract used by focused tests. +func deleteTaskFiles(dataDir, taskID string, explorationID int64) (bool, error) { + stage, err := stageTaskFiles(dataDir, taskID, explorationID) + if err != nil { + return false, err + } + deleted := stage.deleted + if err := stage.commit(); err != nil { + return deleted, err + } + return deleted, nil } func (m *Manager) Task(id string) (*Task, bool) { diff --git a/server/manager_delete_files_test.go b/server/manager_delete_files_test.go new file mode 100644 index 0000000..48de983 --- /dev/null +++ b/server/manager_delete_files_test.go @@ -0,0 +1,280 @@ +package server + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + pgdb "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +func TestSeedAssociatesTargetAssetWithTask(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + task, err := m.CreateTask("seed ownership", "seed ownership", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true}) + }) + + host := fmt.Sprintf("seed-%s.example.test", task.ID) + s := &Server{m: m} + s.seed(task, "https://"+host) + taskID := mustTaskID(t, task.ID) + assets, err := m.Assets().QueryByTask(taskID, "", 20, 0) + if err != nil { + t.Fatal(err) + } + found := false + for _, asset := range assets { + if asset.Domain == host || asset.URL == "https://"+host { + found = true + break + } + } + if !found { + t.Fatalf("seed target was not associated with task %s: %+v", task.ID, assets) + } +} + +func TestDeleteTaskFilesRemovesOnlyOwnedWorkspaceAndTranscripts(t *testing.T) { + dataDir := t.TempDir() + + removed := []struct { + path string + isDir bool + }{ + {path: filepath.Join("tasks", "77"), isDir: true}, + {path: filepath.Join("transcripts", "exp12-main.jsonl")}, + {path: filepath.Join("transcripts", "exp12-planner.jsonl")}, + {path: filepath.Join("transcripts", "exp12-worker-i9.jsonl")}, + {path: filepath.Join("transcripts", "exp12-main"), isDir: true}, + {path: filepath.Join("transcripts", "exp12-worker-i9"), isDir: true}, + } + kept := []struct { + path string + isDir bool + }{ + {path: filepath.Join("tasks", "78"), isDir: true}, + {path: filepath.Join("transcripts", "exp123-main.jsonl")}, + {path: filepath.Join("transcripts", "exp123-main"), isDir: true}, + {path: filepath.Join("transcripts", "exp1-main.jsonl")}, + {path: filepath.Join("transcripts", "exp12main.jsonl")}, + {path: filepath.Join("transcripts", "exp12-main.jsonl.bak")}, + {path: filepath.Join("transcripts", "exp12-main.txt")}, + {path: filepath.Join("transcripts", "conv-12.jsonl")}, + {path: filepath.Join("transcripts", "unrelated-session"), isDir: true}, + } + + for _, fixture := range append(removed, kept...) { + createDeleteFixture(t, dataDir, fixture.path, fixture.isDir) + } + + deleted, err := deleteTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatalf("delete task files: %v", err) + } + if !deleted { + t.Fatal("expected deletion to be reported") + } + for _, fixture := range removed { + assertPathMissing(t, filepath.Join(dataDir, fixture.path)) + } + for _, fixture := range kept { + assertPathExists(t, filepath.Join(dataDir, fixture.path)) + } +} + +func TestDeleteTaskFilesReportsTranscriptOnlyDeletion(t *testing.T) { + dataDir := t.TempDir() + path := filepath.Join("transcripts", "exp44-main.jsonl") + createDeleteFixture(t, dataDir, path, false) + + deleted, err := deleteTaskFiles(dataDir, "missing-task", 44) + if err != nil { + t.Fatalf("delete task files: %v", err) + } + if !deleted { + t.Fatal("expected transcript deletion to be reported") + } + assertPathMissing(t, filepath.Join(dataDir, path)) +} + +func TestDeleteTaskFilesMissingTargetsIsIdempotent(t *testing.T) { + deleted, err := deleteTaskFiles(t.TempDir(), "77", 12) + if err != nil { + t.Fatalf("delete missing task files: %v", err) + } + if deleted { + t.Fatal("missing targets must not be reported as deleted") + } +} + +func TestStageTaskFilesRollbackRestoresWorkspaceAndTranscripts(t *testing.T) { + dataDir := t.TempDir() + paths := []string{ + filepath.Join("tasks", "77", "notes.txt"), + filepath.Join("transcripts", "exp12-main.jsonl"), + filepath.Join("transcripts", "exp12-worker-i9", "sidechain.jsonl"), + } + for _, path := range paths { + createDeleteFixture(t, dataDir, path, false) + } + + stage, err := stageTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatal(err) + } + if !stage.deleted { + t.Fatal("expected files to be staged") + } + for _, path := range paths { + assertPathMissing(t, filepath.Join(dataDir, path)) + } + if err := stage.rollback(); err != nil { + t.Fatal(err) + } + for _, path := range paths { + assertPathExists(t, filepath.Join(dataDir, path)) + } +} + +func TestStageTaskFilesRollbackReportsRestoreFailure(t *testing.T) { + dataDir := t.TempDir() + workspace := filepath.Join(dataDir, "tasks", "77") + createDeleteFixture(t, dataDir, filepath.Join("tasks", "77"), true) + + stage, err := stageTaskFiles(dataDir, "77", 12) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(workspace, []byte("conflict"), 0o644); err != nil { + t.Fatal(err) + } + cause := errors.New("forced PostgreSQL delete failure") + err = rollbackTaskDelete(cause, nil, stage) + if err == nil || !errors.Is(err, cause) || !strings.Contains(err.Error(), "destination already exists") { + t.Fatalf("rollback err=%v, want original and restore errors", err) + } + if _, statErr := os.Stat(stage.stageDir); statErr != nil { + t.Fatalf("staging was removed after failed restore: %v", statErr) + } +} + +func TestManagerDeleteTaskRestoresFilesAndTrafficWhenPostgresDeleteFails(t *testing.T) { + dataDir := t.TempDir() + m, err := NewManager(dataDir, "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + t.Cleanup(func() { _ = m.Close() }) + + task, err := m.CreateTask("delete rollback", "delete rollback", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + taskID := mustTaskID(t, task.ID) + triggerName := fmt.Sprintf("test_fail_task_delete_%d", taskID) + functionName := triggerName + "_fn" + cleanupDB := func() { + _, _ = m.pg.Exec(fmt.Sprintf(`DROP TRIGGER IF EXISTS %s ON tasks`, triggerName)) + _, _ = m.pg.Exec(fmt.Sprintf(`DROP FUNCTION IF EXISTS %s()`, functionName)) + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true, DeleteTraffic: true, DeleteFiles: true}) + } + t.Cleanup(cleanupDB) + + host := fmt.Sprintf("delete-rollback-%d.example.test", taskID) + if _, err := m.Assets().UpsertHTTPService(pgdb.UpsertHTTPServiceReq{URL: "https://" + host, TaskID: taskID}); err != nil { + t.Fatal(err) + } + createDeleteFixture(t, dataDir, filepath.Join("tasks", task.ID), true) + transcriptPath := filepath.Join(dataDir, "transcripts", fmt.Sprintf("exp%d-main.jsonl", task.ExpID)) + createDeleteFixture(t, dataDir, filepath.Join("transcripts", filepath.Base(transcriptPath)), false) + + tr, err := traffic.Open(filepath.Join(dataDir, "traffic"), "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + m.traffic = tr + trafficTree := filepath.Join(dataDir, "traffic", host, "GET", "1-0001") + if err := os.MkdirAll(trafficTree, 0o755); err != nil { + t.Fatal(err) + } + trafficMarker := filepath.Join(trafficTree, "request.http") + if err := os.WriteFile(trafficMarker, []byte("original traffic"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, "1-0001", 1, host, "GET", "/", "https://"+host+"/", 200, "text/plain", 0, 0, host+"/GET/1-0001"); err != nil { + t.Fatal(err) + } + + if _, err := m.pg.Exec(fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $body$ +BEGIN RAISE EXCEPTION 'forced task delete failure'; END $body$`, functionName)); err != nil { + t.Fatal(err) + } + if _, err := m.pg.Exec(fmt.Sprintf(`CREATE TRIGGER %s BEFORE DELETE ON tasks +FOR EACH ROW WHEN (OLD.id = %d) EXECUTE FUNCTION %s()`, triggerName, taskID, functionName)); err != nil { + t.Fatal(err) + } + + _, err = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteTraffic: true, DeleteFiles: true}) + if err == nil || !strings.Contains(err.Error(), "forced task delete failure") { + t.Fatalf("delete err=%v, want injected PostgreSQL failure", err) + } + if got, err := m.pg.GetTask(taskID); err != nil || got == nil { + t.Fatalf("task row was lost after failed delete: task=%+v err=%v", got, err) + } + assertPathExists(t, filepath.Join(dataDir, "tasks", task.ID)) + assertPathExists(t, transcriptPath) + if got, err := os.ReadFile(trafficMarker); err != nil || string(got) != "original traffic" { + t.Fatalf("traffic tree was not restored: content=%q err=%v", got, err) + } + var trafficRows int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&trafficRows); err != nil || trafficRows != 1 { + t.Fatalf("traffic index was not restored: rows=%d err=%v", trafficRows, err) + } +} + +func createDeleteFixture(t *testing.T, root, relative string, isDir bool) { + t.Helper() + path := filepath.Join(root, relative) + if isDir { + if err := os.MkdirAll(filepath.Join(path, "subagents"), 0o755); err != nil { + t.Fatalf("create fixture directory %s: %v", path, err) + } + if err := os.WriteFile(filepath.Join(path, "subagents", "agent-test.jsonl"), []byte("test"), 0o644); err != nil { + t.Fatalf("create fixture sidechain %s: %v", path, err) + } + return + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatalf("create fixture parent %s: %v", path, err) + } + if err := os.WriteFile(path, []byte("test"), 0o644); err != nil { + t.Fatalf("create fixture file %s: %v", path, err) + } +} + +func assertPathMissing(t *testing.T, path string) { + t.Helper() + if _, err := os.Lstat(path); !os.IsNotExist(err) { + t.Fatalf("expected %s to be removed, stat err=%v", path, err) + } +} + +func assertPathExists(t *testing.T, path string) { + t.Helper() + if _, err := os.Lstat(path); err != nil { + t.Fatalf("expected %s to remain: %v", path, err) + } +} diff --git a/server/orchestration.go b/server/orchestration.go index 011b0f6..9ed4446 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -199,12 +199,13 @@ func (s *Server) toolListTasks() actool.CoreTool { if t.ParentRef != "" { row["parent_ref"] = t.ParentRef } - if t.LLMProfileID == nil { + llmState := t.llmStateSnapshot() + if llmState.ProfileID == nil { row["llm_profile"] = "(激活配置)" - } else if n, ok := profName[*t.LLMProfileID]; ok { + } else if n, ok := profName[*llmState.ProfileID]; ok { row["llm_profile"] = n } else { - row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID) + row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *llmState.ProfileID) } out = append(out, row) } @@ -305,6 +306,7 @@ func (s *Server) toolPauseTask() actool.CoreTool { return actool.Errorf("task 不存在: " + a.TaskID), nil } s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator) + s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask) return actool.Text("task paused: " + a.TaskID), nil }) } @@ -419,7 +421,7 @@ func (s *Server) deriveTaskStatus(t *Task) string { return t.Status case t.Paused || s.engine.IsPaused(t.ID): return "paused" - case s.engine.Ready() && s.engine.Started(t.ID): + case s.engine.ReadyFor(t) && s.engine.Started(t.ID): return "running" } return "created" diff --git a/server/server.go b/server/server.go index 2fe0d1d..2f44d26 100644 --- a/server/server.go +++ b/server/server.go @@ -12,6 +12,7 @@ import ( "net/url" "path/filepath" "regexp" + "sort" "strconv" "strings" "sync" @@ -51,14 +52,12 @@ type Server struct { cfgMu sync.Mutex mainAgent *agent.MainAgent // nil when no LLM provider is configured chatAgent *agent.ChatAgent // conversational runner for the chat page; nil w/o LLM - llmCfg agent.Config // current LLM config (key not exposed) - llmOn bool - llmProf string // active LLM profile name (for llmrec tagging) // llmProv is the fully-decorated global provider (recorder + failover chain) - // applyLLM installed. Kept so one-off callers (goal decomposition) reuse the - // same instance instead of building a private one that skips the rate limiter, - // the recorder and failover. + // installed by applyLLM. Task routers use it after an explicit chain is cleared. llmProv llm.Provider + llmCfg agent.Config // current LLM config (key not exposed) + llmOn bool + llmProf string // active LLM profile name (for llmrec tagging) // chatBusy guards the per-task main-agent run: the chat handler launches the // agent on the server's background ctx (not the request ctx) and returns @@ -97,18 +96,25 @@ type Server struct { // provByProfile caches ONE provider per LLM profile id so every agent bound or // pinned to the same profile shares a single provider instance — hence one rate - // limiter. Separate from profAgents (whole planner/worker pairs). The globally-active - // provider (built in applyLLM) lives OUTSIDE this cache, so binding an agent to the - // currently-active profile builds a second instance — harmless (that binding equals - // leaving it unset), just not deduped. Cleared alongside profAgents on profile edits. + // limiter. Separate from profAgents (whole planner/worker pairs). applyLLM also + // resolves the persisted global-active profile through this cache, so global + // fallback and task chains do not accidentally double the configured request rate. + // Cleared alongside profAgents on profile edits, then repopulated by active reapply. provCacheMu sync.Mutex provByProfile map[int64]*provEntry + provCacheGen uint64 // llmHealth is the process-wide circuit-breaker state for LLM failover (轮询). // It deliberately lives OUTSIDE the provider caches: rebuilding the chain // (saving an unrelated profile, flipping a setting) must not erase what we // learned about which backends are out of credit / rate-limited. llmHealth *llmpool.Registry + + // Explicit task chains use one task-scoped dynamic provider shared by goals, + // planner, workers, and the main agent. Bundles are stable; their runtime reads + // the persisted current profile at every new LLM call. + taskAgentMu sync.Mutex + taskAgents map[string]*taskAgentBundle } // profBundle is a planner/worker pair built from one LLM profile. @@ -143,14 +149,17 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi chatCancel: map[string]context.CancelCauseFunc{}, triggerQ: map[string][]triggeredRun{}, triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{}, profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{}, - provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg)} - // per-task LLM: a task pinned to a specific profile runs on that profile's - // dedicated planner/worker; unpinned tasks fall back to the global active pair. - s.engine.SetAgentResolver(func(t *Task) (*agent.Planner, *agent.Worker) { - if t.LLMProfileID == nil { + provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg), + taskAgents: map[string]*taskAgentBundle{}} + // Every task uses a stable task router. An empty explicit chain is resolved by + // that router through Agent bindings and then the global provider, so adding a + // first chain to a running task takes effect on its very next LLM call. + s.engine.SetAuthoritativeAgentResolver(func(t *Task) (*agent.Planner, *agent.Worker) { + if !s.taskRuntimeAvailable(t, "planner", "worker") { return nil, nil } - return s.agentsForProfile(*t.LLMProfileID) + b := s.agentsForTask(t) + return b.pl, b.wk }) // Wire DB-stored prompt templates into the agents (新版方案 §3.3 / §5a). With no // override row, agents keep their built-in defaults — behavior is unchanged. @@ -297,7 +306,7 @@ func (s *Server) loadLLMConfig() (agent.Config, bool) { } // saveLLMConfig persists the LLM config as the active "default" profile in PG. -func (s *Server) saveLLMConfig(cfg agent.Config) { +func (s *Server) saveLLMConfig(cfg agent.Config) error { format := cfg.Provider() if format != "anthropic" { format = "openai" @@ -322,9 +331,10 @@ func (s *Server) saveLLMConfig(cfg agent.Config) { ContextWindowK: cfg.ContextWindowK, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, Priority: priority, PoolExclude: poolExclude, }) - if err == nil { - _ = s.m.pg.SetActiveProfile(newID) + if err != nil { + return err } + return s.m.pg.SetActiveProfile(newID) } // reapplyActiveProfile hot-reloads the engine from the active DB profile so that @@ -380,15 +390,26 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent // applyLLM (re)builds the planner/worker/main-agent from cfg and installs them on // the running engine as the GLOBAL active pair. Safe to call at runtime (UI configures LLM). func (s *Server) applyLLM(cfg agent.Config) error { - prov, err := cfg.NewProvider() - if err != nil { - return err + var prov llm.Provider + // A persisted active profile must use the same cached provider as task chains + // and Agent bindings, otherwise each path owns a separate rate limiter. + if active, _ := s.m.pg.ActiveProfile(); active != nil { + if activeCfg, ok := s.loadProfileConfig(active.ID); ok && activeCfg == cfg { + prov, _, _ = s.providerForProfile(active.ID) + } + } + if prov == nil { + var err error + prov, err = cfg.NewProvider() + if err != nil { + return err + } + // Non-persisted/env configs do not have a profile cache key. + s.cfgMu.Lock() + profName := s.llmProf + s.cfgMu.Unlock() + prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) } - // Wrap provider with LLM call recorder (persists request/response to PG). - s.cfgMu.Lock() - profName := s.llmProf - s.cfgMu.Unlock() - prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) // LLM 轮询(默认关):把激活配置包进故障转移链,当前配置不可用时自动切下一个。 // 只影响「走全局激活配置」的这条路径——agent 绑定 / 任务 pin 的走 providerForProfile, // 默认仍然独占该配置(见 poolForBinding)。关闭或无备选时返回原 provider,行为不变。 @@ -414,10 +435,11 @@ func (s *Server) applyLLM(cfg agent.Config) error { s.chatAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) s.chatAgent.SetWebSearch(s.m.WebSearchOpts()) s.chatAgent.SetGuard(s.chatGuard()) - s.llmCfg = cfg s.llmProv = prov + s.llmCfg = cfg s.llmOn = true s.cfgMu.Unlock() + s.invalidateTaskAgents() // wake the active task so a task created while idle starts exploring. if t := s.m.ActiveTask(); t != nil { @@ -466,6 +488,7 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) s.provCacheMu.Unlock() return e.prov, e.cfg, true } + generation := s.provCacheGen s.provCacheMu.Unlock() cfg, ok := s.loadProfileConfig(id) if !ok { @@ -481,6 +504,10 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, s.m.LLMRecordEnabled) } s.provCacheMu.Lock() + if generation != s.provCacheGen { + s.provCacheMu.Unlock() + return s.providerForProfile(id) + } if e := s.provByProfile[id]; e != nil { // lost the race → keep the winner prov, cfg = e.prov, e.cfg } else { @@ -571,8 +598,10 @@ func (s *Server) invalidateProfileAgents() { s.profChatAgents = map[int64]*agent.ChatAgent{} s.profMu.Unlock() s.provCacheMu.Lock() + s.provCacheGen++ s.provByProfile = map[int64]*provEntry{} s.provCacheMu.Unlock() + s.invalidateTaskAgents() } func (s *Server) mainAgentRef() *agent.MainAgent { @@ -619,6 +648,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /api/workspace/upload", s.wsUpload) mux.HandleFunc("GET /api/tasks/{id}/scope", s.taskScopeList) mux.HandleFunc("POST /api/tasks/{id}/control", s.control) + mux.HandleFunc("PUT /api/tasks/{id}/llm", s.updateTaskLLMProfiles) + mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/control", s.controlIntent) mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/rerun", s.rerunIntent) // 重跑单条 blocked/exhausted/stopped 意图 mux.HandleFunc("POST /api/tasks/{id}/intents/rerun-blocked", s.rerunBlocked) // 批量重跑本任务全部 blocked 意图 mux.HandleFunc("POST /api/active", s.setActive) @@ -821,6 +852,7 @@ func (s *Server) stats(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, out) // no task selected: only global fields return } + out["llm_configured"] = s.engine.ReadyFor(t) st, _ := t.Store.Stats() out["exploration"] = st @@ -842,27 +874,31 @@ func (s *Server) stats(w http.ResponseWriter, r *http.Request) { } last := s.engine.LastActivity(t.ID) paused := s.engine.IsPaused(t.ID) - running := s.engine.Ready() && s.engine.Started(t.ID) && !paused - // stalled: running but no activity for a while and nothing in flight. - stalled := running && inFlight == 0 && last > 0 && time.Now().Unix()-last > 120 + activeCalls := s.engine.ActiveLLMCalls(t.ID) + running := s.engine.ReadyFor(t) && s.engine.Started(t.ID) && !paused + // A live event-driven engine with no work is healthy idle. Only a persisted + // running intent without any corresponding LLM call can be considered stalled. + stalled := running && activeCalls == 0 && inFlight > 0 && last > 0 && time.Now().Unix()-last > 120 // engine_mode follows the spec enum (exploring|paused|stalled|idle). + engineMode := "idle" switch { case paused: - out["engine_mode"] = "paused" + engineMode = "paused" case stalled: - out["engine_mode"] = "stalled" - case running: - out["engine_mode"] = "exploring" - default: - out["engine_mode"] = "idle" + engineMode = "stalled" + case running && activeCalls > 0: + engineMode = "exploring" } + out["engine_mode"] = engineMode out["active_task"] = map[string]any{ "id": t.ID, "description": t.Description, "goal": t.Goal, "running": running, "paused": paused, + "engine_mode": engineMode, "in_flight": inFlight, + "llm_in_flight": activeCalls, "last_activity": last, "stalled": stalled, "goals_total": len(goals), @@ -882,18 +918,7 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { goalCounts, _ := s.m.PG().GoalCountsAll() // goal progress per exploration, one query dtos := make([]TaskDTO, 0, len(list)) for _, t := range list { - status := "created" - switch { - case isTerminalStatus(t.Status): // 持久化终态优先（done/failed/timeout） - status = t.Status - case t.Queued: // 因并发上限挂起、等待空位(尚未启动) - status = "queued" - case t.Paused || s.engine.IsPaused(t.ID): - status = "paused" - case s.engine.Ready() && s.engine.Started(t.ID): - status = "running" - } - dto := taskDTO(t, status) + dto := taskDTO(t, s.resolvedTaskStatus(t)) dto.Tokens = tokenTotalDTO(toks[t.ExpID]) // prefer the live in-memory heartbeat (fresher) and fall back to the // persisted max activity time (survives restarts) for run-duration display. @@ -909,6 +934,28 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"tasks": dtos, "active": active}) } +func (s *Server) resolvedTaskStatus(t *Task) string { + if t == nil { + return "created" + } + switch { + case isTerminalStatus(t.Status): + return t.Status + case t.Queued: + return "queued" + case t.Paused || s.engine.IsPaused(t.ID): + return "paused" + case t.llmStateSnapshot().FailoverState == "chain_exhausted" && s.engine.Started(t.ID): + // LLM readiness is an execution dependency, not lifecycle state. Keeping + // an exhausted task running lets the user edit/reset its chain directly. + return "running" + case s.engine.ReadyFor(t) && s.engine.Started(t.ID): + return "running" + default: + return "created" + } +} + func (s *Server) setActive(w http.ResponseWriter, r *http.Request) { var req struct { ID string `json:"id"` @@ -946,6 +993,9 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { switch req.Action { case "pause": s.engine.Pause(t.ID, agent.AbortPausedByUser) + // Main Agent turns run outside Engine's planner/worker context so they need + // their own cancellation. New turns are rejected while the task is paused. + s.cancelTaskChat(t.ID, agent.AbortChatPausedWithTask) case "resume": s.engine.Run(s.ctx, t) // ensure loops are alive, then un-pause + nudge s.engine.Resume(t) @@ -962,6 +1012,99 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"id": t.ID, "paused": paused}) } +// controlIntent pauses, resumes, or cancels one local worker intent. A running +// worker is always stopped before state mutation/cleanup, preventing tool output +// that arrives after the user action from recreating deleted blackboard records. +func (s *Server) controlIntent(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + if !s.engine.beginTaskOperation(t.ID) { + writeErr(w, 409, "任务正在删除，无法控制意图") + return + } + defer s.engine.decInflight(t.ID) + + iid, err := strconv.ParseInt(r.PathValue("iid"), 10, 64) + if err != nil || iid <= 0 { + writeErr(w, 400, "bad intent id") + return + } + var req struct { + Action string `json:"action"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeErr(w, 400, "bad json: "+err.Error()) + return + } + if req.Action != "pause" && req.Action != "resume" && req.Action != "cancel" { + writeErr(w, 400, "action must be pause|resume|cancel") + return + } + + node, err := t.Store.GetNode(iid) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if node == nil { + if inherited, sourceErr := t.Store.GetNodeWithSources(iid); sourceErr == nil && inherited != nil && inherited.Inherited { + writeErr(w, 409, "继承意图为只读，不能控制") + } else { + writeErr(w, 404, "intent not found") + } + return + } + if node.Kind != db.KindIntent { + writeErr(w, 409, "node is not an intent") + return + } + + switch req.Action { + case "pause": + if node.State != "running" { + writeErr(w, 409, "仅运行中的意图可以暂停") + return + } + if err := s.engine.ControlWork(iid, "pause"); err != nil { + writeErr(w, 409, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"id": iid, "state": "paused"}) + case "resume": + if node.State != "paused" { + writeErr(w, 409, "仅已暂停的意图可以恢复") + return + } + if err := t.Store.SetIntentState(iid, "open"); err != nil { + writeErr(w, 500, err.Error()) + return + } + t.Notify() + writeJSON(w, 200, map[string]any{"id": iid, "state": "open"}) + case "cancel": + if node.State != "running" && node.State != "paused" { + writeErr(w, 409, "仅运行中或已暂停的意图可以取消") + return + } + if node.State == "running" { + if err := s.engine.ControlWork(iid, "cancel"); err != nil { + writeErr(w, 409, err.Error()) + return + } + } + deleted, err := t.Store.CancelIntent(iid) + if err != nil { + writeErr(w, 409, err.Error()) + return + } + t.Notify() + writeJSON(w, 200, map[string]any{"id": iid, "state": "cancelled", "deleted": deleted}) + } +} + // rerunIntent 重跑一条没跑成功的意图(blocked/exhausted/stopped):把它置回 open,worker // 会重新认领、从头再跑(已写回图谱的 fact/finding/asset 保留);若任务已终态/暂停则顺带复活。 // 用于「出错的 work 点击继续运行」——网络/LLM 抖动导致 blocked 后可一键重试。 @@ -1063,11 +1206,24 @@ func (s *Server) setLLM(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "api_key required") return } + // Validate provider construction before persisting it as the active profile. + if _, err := cfg.NewProvider(); err != nil { + writeErr(w, 400, "provider init failed: "+err.Error()) + return + } + if err := s.saveLLMConfig(cfg); err != nil { + writeErr(w, 500, "persist provider failed: "+err.Error()) + return + } + s.invalidateProfileAgents() + if _, ok := s.loadLLMConfig(); !ok { + writeErr(w, 500, "saved provider is unavailable") + return + } if err := s.applyLLM(cfg); err != nil { writeErr(w, 400, "provider init failed: "+err.Error()) return } - s.saveLLMConfig(cfg) // persist to DB log.Printf("[engine] LLM configured via UI: %s / %s", cfg.Provider(), cfg.Model) s.getLLM(w, r) } @@ -1116,12 +1272,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { } type createTaskReq struct { - Description string `json:"description"` - Goal string `json:"goal"` - LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 - TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时 - PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600 - SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。 + Description string `json:"description"` + Goal string `json:"goal"` + LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 + LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效 + SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务 + TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时 + PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600 + SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。 } func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { @@ -1133,18 +1291,39 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { if strings.TrimSpace(req.Description) == "" { req.Description = "未命名任务" } - // validate the pinned profile up front (must exist + be usable), so a bad id fails - // task creation instead of silently falling back to the active profile at run time. - if req.LLMProfileID != nil { - if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok { - writeErr(w, 400, "所选 LLM 配置不存在或未设置 API Key") - return - } + if len(req.LLMProfileIDs) == 0 && req.LLMProfileID != nil { + req.LLMProfileIDs = []int64{*req.LLMProfileID} + } + if err := s.validateTaskProfileIDs(req.LLMProfileIDs); err != nil { + writeErr(w, 400, err.Error()) + return } if req.TimeoutSeconds < 0 { req.TimeoutSeconds = 0 } - t, err := s.m.CreateTask(req.Description, req.Goal, req.LLMProfileID, req.TimeoutSeconds, req.PlanHeartbeatSeconds) + if len(req.SourceTaskIDs) > db.MaxTaskSourceCount { + writeErr(w, 400, fmt.Sprintf("关联任务最多选择 %d 个", db.MaxTaskSourceCount)) + return + } + sourceIDs := make([]int64, 0, len(req.SourceTaskIDs)) + seenSources := map[int64]bool{} + for _, raw := range req.SourceTaskIDs { + id, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64) + if err != nil || id <= 0 || seenSources[id] { + writeErr(w, 400, "关联任务 id 无效或重复") + return + } + if _, ok := s.m.Task(strconv.FormatInt(id, 10)); !ok { + writeErr(w, 400, fmt.Sprintf("关联任务 #%d 不存在", id)) + return + } + seenSources[id] = true + sourceIDs = append(sourceIDs, id) + } + t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{ + SourceTaskIDs: sourceIDs, LLMProfileIDs: req.LLMProfileIDs, + TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds, + }) if err != nil { writeErr(w, 500, err.Error()) return @@ -1153,7 +1332,73 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { // 共享的建后流程(seed + 种子意图 + 后台目标分解 + engine.Run),与 spawn_task 复用同一段。 // launchTask 内部异步,不阻塞 UI —— 目标分解在后台可见地进行。 s.launchTask(t, req.Description+" "+req.Goal, req.SeedFirstIntent != nil && *req.SeedFirstIntent) - writeJSON(w, 201, t) + writeJSON(w, 201, taskDTO(t, s.resolvedTaskStatus(t))) +} + +func (s *Server) validateTaskProfileIDs(ids []int64) error { + seen := map[int64]bool{} + for _, id := range ids { + if id <= 0 || seen[id] { + return fmt.Errorf("LLM 配置 id 无效或重复") + } + seen[id] = true + if _, ok := s.loadProfileConfig(id); !ok { + return fmt.Errorf("LLM 配置 #%d 不存在或未设置 API Key", id) + } + } + return nil +} + +func (s *Server) updateTaskLLMProfiles(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + if isTerminalStatus(t.Status) { + writeErr(w, 409, "已结束的任务不能修改 LLM 配置") + return + } + before := t.llmStateSnapshot() + var req struct { + LLMProfileIDs []int64 `json:"llm_profile_ids"` + ActiveLLMProfileID *int64 `json:"active_llm_profile_id"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeErr(w, 400, "bad json: "+err.Error()) + return + } + if err := s.validateTaskProfileIDs(req.LLMProfileIDs); err != nil { + writeErr(w, 400, err.Error()) + return + } + active := int64(0) + if req.ActiveLLMProfileID != nil { + active = *req.ActiveLLMProfileID + } + reopened, err := s.m.ReplaceTaskLLMProfiles(t.ID, req.LLMProfileIDs, active) + if err != nil { + if strings.Contains(err.Error(), "terminal") { + writeErr(w, 409, err.Error()) + } else { + writeErr(w, 400, err.Error()) + } + return + } + // Profile edits affect only subsequent LLM calls. Existing in-flight calls + // retain their provider; a nudge lets an idle running task resume promptly. + t.Notify() + llmState := t.llmStateSnapshot() + result := map[string]any{ + "id": t.ID, "llm_profile_ids": llmState.ProfileIDs, + "active_llm_profile_id": llmState.ActiveID, + "llm_failover_state": llmState.FailoverState, "reopened_intents": reopened, + } + if !sameOptionalID(before.ActiveID, llmState.ActiveID) { + event := s.emitManualTaskLLMSwitch(t, before.ActiveID, llmState.ActiveID) + result["switch_event"] = activityDTO(event) + } + writeJSON(w, 200, result) } var ( @@ -1241,12 +1486,13 @@ func (s *Server) seed(t *Task, text string) { } var rootID int64 if as := s.m.Assets(); as != nil { + taskID, _ := strconv.ParseInt(t.ID, 10, 64) if net.ParseIP(host) != nil { - rootID, _ = as.UpsertIP(db.UpsertIPReq{IP: host}) + rootID, _ = as.UpsertIP(db.UpsertIPReq{IP: host, TaskID: taskID}) } else if scheme == "https" || scheme == "http" { - rootID, _ = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: u}) + rootID, _ = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: u, TaskID: taskID}) } else { - rootID, _ = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host}) + rootID, _ = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: taskID}) } } // anchor the seeded assets to this task's begin root as lineage/provenance @@ -1286,7 +1532,7 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "task not found") return } - writeJSON(w, 200, t) + writeJSON(w, 200, taskDTO(t, s.resolvedTaskStatus(t))) } // taskCoverage returns a task's rough asset test coverage (denominator/tested/backlog). @@ -1302,7 +1548,7 @@ func (s *Server) taskCoverage(w http.ResponseWriter, r *http.Request) { return } taskID, _ := strconv.ParseInt(t.ID, 10, 64) - cov, err := as.TaskCoverage(taskID, t.ExpID) + cov, err := as.TaskCoverageWithSources(taskID) if err != nil { writeErr(w, 500, err.Error()) return @@ -1345,22 +1591,23 @@ func (s *Server) taskAssetRefs(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "需要 asset_id") return } - refs, err := t.Store.AssetRefs(assetID) + refs, err := t.Store.AssetRefsWithSources(assetID) if err != nil { writeErr(w, 500, err.Error()) return } - intents := []db.AssetRef{} - facts := []db.AssetRef{} - findings := []db.AssetRef{} + intents := []CoverageAssetRefDTO{} + facts := []CoverageAssetRefDTO{} + findings := []CoverageAssetRefDTO{} for _, ref := range refs { + dto := coverageAssetRefDTO(ref) switch ref.Kind { case "intent": - intents = append(intents, ref) + intents = append(intents, dto) case "fact": - facts = append(facts, ref) + facts = append(facts, dto) case "finding": - findings = append(findings, ref) + findings = append(findings, dto) } } writeJSON(w, 200, map[string]any{"intents": intents, "facts": facts, "findings": findings}) @@ -1379,7 +1626,7 @@ func (s *Server) taskScopeList(w http.ResponseWriter, r *http.Request) { return } taskID, _ := strconv.ParseInt(t.ID, 10, 64) - rows, err := as.ListTaskScope(taskID) + rows, err := as.ListTaskScopeWithSources(taskID) if err != nil { writeErr(w, 500, err.Error()) return @@ -1443,7 +1690,11 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, []any{}) return } - f, _ := t.Store.ListByKind(db.KindFinding, 200) + f, err := t.Store.ListByKind(db.KindFinding, 200) + if err != nil { + writeErr(w, 500, err.Error()) + return + } tid, _ := strconv.ParseInt(t.ID, 10, 64) meta, err := s.m.pg.FindingMetaByNodeID(tid) if err != nil { @@ -1454,7 +1705,42 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) { aidSet = append(aidSet, m.AssetIDs...) } assets := s.resolveAssetIDs(aidSet) - writeJSON(w, 200, findingDTOsForTask(t, f, meta, assets)) + out := findingDTOsForTask(t, f, meta, assets) + + // Related-task findings are a live, read-only view. Provenance metadata lets + // the task UI suppress mutation affordances while retaining stable ids. + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + nodes, listErr := source.Store.ListByKind(db.KindFinding, 200) + if listErr != nil { + writeErr(w, 500, listErr.Error()) + return + } + for _, node := range nodes { + node.SourceTaskID = source.Task.TaskID + node.Inherited = true + } + sourceMeta, metaErr := s.m.pg.FindingMetaByNodeID(source.Task.TaskID) + if metaErr != nil { + log.Printf("[findings] source_task=%d meta: %v", source.Task.TaskID, metaErr) + } + var sourceAssetIDs []int64 + for _, item := range sourceMeta { + sourceAssetIDs = append(sourceAssetIDs, item.AssetIDs...) + } + out = append(out, findingDTOsForOwner( + i64s(source.Task.TaskID), + source.Task.Description, + nodes, + sourceMeta, + s.resolveAssetIDs(sourceAssetIDs), + )...) + } + writeJSON(w, 200, out) } // normFilter maps the frontend's "all" sentinel (and empty) to "" so the DB layer @@ -1529,7 +1815,41 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) { return } assets := s.resolveAssetIDs(f.AssetIDs) - writeJSON(w, 200, findingFromDB(f, assets)) + dto := findingFromDB(f, assets) + if contextTaskID := strings.TrimSpace(r.URL.Query().Get("context_task")); contextTaskID != "" { + contextTask := s.m.ResolveTask(contextTaskID) + if contextTask == nil { + writeErr(w, 404, "context task not found") + return + } + sourceTaskID, inherited, allowed := findingProvenanceInTask(contextTask, f.TaskID) + if !allowed { + writeErr(w, 404, "finding not available in task context") + return + } + dto.SourceTaskID = sourceTaskID + dto.Inherited = inherited + } + writeJSON(w, 200, dto) +} + +func findingProvenanceInTask(contextTask *Task, findingTaskID *int64) (sourceTaskID string, inherited, allowed bool) { + if contextTask == nil || findingTaskID == nil { + return "", false, false + } + contextID, err := strconv.ParseInt(contextTask.ID, 10, 64) + if err != nil { + return "", false, false + } + if *findingTaskID == contextID { + return "", false, true + } + for _, sourceID := range contextTask.SourceTaskIDs { + if sourceID == *findingTaskID { + return i64s(sourceID), true, true + } + } + return "", false, false } // findingLineage returns the exploration sub-DAG from the task root to this @@ -1699,17 +2019,108 @@ func (s *Server) intents(w http.ResponseWriter, r *http.Request) { return } // Paged form: ?before= (or ?page as a marker) → {items, has_more} so the - // worker session list can reach past the old fixed 300 boundary on scroll. + // worker session list can reach past the old fixed 300 boundary on scroll. It + // also includes immutable intent results from directly related tasks; those + // rows never participate in this task's frontier or claim path. before := int64(atoiDefault(q.Get("before"), 0)) - in, hasMore, err := t.Store.ListByKindPage(db.KindIntent, before, limit) + in, hasMore, err := taskIntentHistoryPage(t.Store, before, limit, false, 0) if err != nil { log.Printf("[intents] task=%s before=%d limit=%d: %v", t.ID, before, limit, err) writeErr(w, 500, err.Error()) return } + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + items, sourceMore, sourceErr := taskIntentHistoryPage(source.Store, before, limit, true, source.Task.TaskID) + if sourceErr != nil { + writeErr(w, 500, sourceErr.Error()) + return + } + in = append(in, items...) + hasMore = hasMore || sourceMore + } + sort.Slice(in, func(i, j int) bool { return in[i].ID > in[j].ID }) + if len(in) > limit { + in = in[:limit] + hasMore = true + } writeJSON(w, 200, map[string]any{"items": taskNodeDTOs(in), "has_more": hasMore}) } +func inheritedIntentResult(state string) bool { + switch state { + case "done", "blocked", "exhausted", "stopped": + return true + default: + return false + } +} + +// inheritedGraphSnapshot exposes immutable source-task history without leaking +// work that is still open or running. Edges are filtered with the nodes so the +// response cannot contain dangling ids that reveal hidden live intent topology. +func inheritedGraphSnapshot(nodes []*db.Node, edges []db.Edge, sourceTaskID int64) ([]*db.Node, []db.Edge) { + visible := make(map[int64]struct{}, len(nodes)) + filteredNodes := make([]*db.Node, 0, len(nodes)) + for _, node := range nodes { + if node == nil || (node.Kind == db.KindIntent && !inheritedIntentResult(node.State)) { + continue + } + node.SourceTaskID = sourceTaskID + node.Inherited = true + visible[node.ID] = struct{}{} + filteredNodes = append(filteredNodes, node) + } + filteredEdges := make([]db.Edge, 0, len(edges)) + for _, edge := range edges { + if _, ok := visible[edge.From]; !ok { + continue + } + if _, ok := visible[edge.To]; !ok { + continue + } + filteredEdges = append(filteredEdges, edge) + } + return filteredNodes, filteredEdges +} + +// taskIntentHistoryPage returns one newest-first page for an exploration. +// A source may have many open/running intents, so keep paging until enough +// historical results are collected instead of leaking its frontier into the UI. +func taskIntentHistoryPage(store *db.ExplorationStore, before int64, limit int, inherited bool, sourceTaskID int64) ([]*db.Node, bool, error) { + if !inherited { + return store.ListByKindPage(db.KindIntent, before, limit) + } + batch := max(limit, 300) + cursor := before + out := make([]*db.Node, 0, limit+1) + for { + page, more, err := store.ListByKindPage(db.KindIntent, cursor, batch) + if err != nil { + return nil, false, err + } + for _, node := range page { + if !inheritedIntentResult(node.State) { + continue + } + node.SourceTaskID = sourceTaskID + node.Inherited = true + out = append(out, node) + if len(out) > limit { + return out[:limit], true, nil + } + } + if !more || len(page) == 0 { + return out, false, nil + } + cursor = page[len(page)-1].ID + } +} + // explorationGraph returns the whole exploration chain (task graph) as nodes+edges. func (s *Server) explorationGraph(w http.ResponseWriter, r *http.Request) { t := s.m.ResolveTask(r.URL.Query().Get("task")) @@ -1717,8 +2128,36 @@ func (s *Server) explorationGraph(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"nodes": []any{}, "edges": []any{}}) return } - nodes, _ := t.Store.Nodes(2000) - edges, _ := t.Store.Edges(5000) + nodes, err := t.Store.Nodes(2000) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + edges, err := t.Store.Edges(5000) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + sources, err := t.Store.DirectSourceStores() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + for _, source := range sources { + sourceNodes, nodeErr := source.Store.Nodes(2000) + if nodeErr != nil { + writeErr(w, 500, nodeErr.Error()) + return + } + sourceEdges, edgeErr := source.Store.Edges(5000) + if edgeErr != nil { + writeErr(w, 500, edgeErr.Error()) + return + } + sourceNodes, sourceEdges = inheritedGraphSnapshot(sourceNodes, sourceEdges, source.Task.TaskID) + nodes = append(nodes, sourceNodes...) + edges = append(edges, sourceEdges...) + } writeJSON(w, 200, map[string]any{"nodes": taskNodeDTOs(nodes), "edges": edgeDTOs(edges)}) } @@ -1737,7 +2176,16 @@ func (s *Server) activity(w http.ResponseWriter, r *http.Request) { intentPtr = &n } } - items, cursor, err := t.Store.ActivityList(intentPtr, since, limit) + var ( + items []db.Activity + cursor int64 + err error + ) + if intentPtr != nil { + items, cursor, err = t.Store.ActivityListWithSources(*intentPtr, since, limit) + } else { + items, cursor, err = t.Store.ActivityList(nil, since, limit) + } if err != nil { log.Printf("[activity] task=%s since=%d limit=%d intent=%v: %v", t.ID, since, limit, intentPtr, err) writeErr(w, 500, err.Error()) @@ -1765,6 +2213,38 @@ func parseActivitySession(sess string) (db.ActivitySessionFilter, bool) { return db.ActivitySessionFilter{}, false } +// activitySessionStore resolves a worker session against the current task and +// its direct sources. Planner/main always stay local; inherited intent sessions +// are immutable history and use the source exploration only for reads. +func activitySessionStore(t *Task, filter db.ActivitySessionFilter) (*db.ExplorationStore, int64, error) { + if filter.NodeID == nil { + return t.Store, 0, nil + } + node, err := t.Store.GetNodeWithSources(*filter.NodeID) + if err != nil { + return nil, 0, err + } + if node == nil || node.Kind != db.KindIntent { + return nil, 0, nil + } + if !node.Inherited { + return t.Store, 0, nil + } + if !inheritedIntentResult(node.State) { + return nil, 0, nil + } + sources, err := t.Store.DirectSourceStores() + if err != nil { + return nil, 0, err + } + for _, source := range sources { + if source.Task.TaskID == node.SourceTaskID { + return source.Store, source.Task.TaskID, nil + } + } + return nil, 0, nil +} + // activityHistory serves one reverse-paginated page of a session's activity history. // The latest page (no ?before) opens a session; ?before= pulls the older page on // scroll-up. snapshot_cursor is the TASK-level max id at query time — the client uses @@ -1785,18 +2265,43 @@ func (s *Server) activityHistory(w http.ResponseWriter, r *http.Request) { } before := int64(atoiDefault(q.Get("before"), 0)) limit := min(atoiDefault(q.Get("limit"), 200), 500) // cap so one request can't pull an unbounded slice + store, sourceTaskID, err := activitySessionStore(t, filter) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if store == nil { + writeErr(w, 404, "session not found") + return + } + // The browser tails only this task's broadcaster. Keep the SSE join cursor + // local even when the opened worker transcript comes from a related task. snapshot, err := t.Store.ActivityMaxID() if err != nil { log.Printf("[activity/history] task=%s session=%s snapshot: %v", t.ID, sess, err) writeErr(w, 500, err.Error()) return } - items, hasMore, err := t.Store.ActivityPage(filter, before, limit) + var items []db.Activity + var hasMore bool + if sourceTaskID > 0 { + // Source sessions are readable only while the intent remains terminal. The + // DB query also removes model reasoning/accounting rows from inherited data. + items, hasMore, err = store.ActivityPageForTerminalIntent(*filter.NodeID, before, limit) + } else { + items, hasMore, err = store.ActivityPage(filter, before, limit) + } if err != nil { log.Printf("[activity/history] task=%s session=%s before=%d limit=%d: %v", t.ID, sess, before, limit, err) writeErr(w, 500, err.Error()) return } + if sourceTaskID > 0 { + for i := range items { + items[i].SourceTaskID = sourceTaskID + items[i].Inherited = true + } + } earliest := before if len(items) > 0 { earliest = items[0].ID @@ -2072,7 +2577,7 @@ func (s *Server) activityDetail(w http.ResponseWriter, r *http.Request) { return } seq, _ := strconv.ParseInt(r.PathValue("seq"), 10, 64) - d, err := t.Store.ActivityDetail(seq) + d, err := taskActivityDetail(t, seq) if err != nil { writeErr(w, 500, err.Error()) return @@ -2080,6 +2585,14 @@ func (s *Server) activityDetail(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"detail": d}) } +// taskActivityDetail keeps the legacy local-task behavior (including thinking +// rows), while inherited details are restricted to terminal worker intents. +// This prevents a guessed global activity id from exposing source planner/main +// or in-flight transcripts. +func taskActivityDetail(t *Task, seq int64) (string, error) { + return t.Store.ActivityDetailWithSources(seq) +} + func (s *Server) getTraffic(w http.ResponseWriter, r *http.Request) { tr := s.m.Traffic() if tr == nil { @@ -2423,6 +2936,14 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "no active task") return } + if s.engine.IsDeleting(t.ID) { + writeErr(w, 409, "任务正在删除，无法发送新消息") + return + } + if s.engine.IsPaused(t.ID) { + writeErr(w, 409, "任务已暂停，无法发送新消息") + return + } var req struct { Message string `json:"message"` Attachments []chatAttachment `json:"attachments,omitempty"` // 方式1 上传的文件(路径相对任务工作目录) @@ -2431,29 +2952,36 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, err.Error()) return } + // Admission is serialized with deletion's chat cancellation. Mark every chat + // turn busy before its first activity/file write, including rule-mode turns. + // If deletion wins the race, the second barrier check rejects this request. + s.chatMu.Lock() + if s.engine.IsDeleting(t.ID) || s.engine.IsPaused(t.ID) { + s.chatMu.Unlock() + writeErr(w, 409, "任务正在删除或已暂停，无法发送新消息") + return + } + if s.chatBusy[t.ID] { + s.chatMu.Unlock() + writeErr(w, 409, "主 Agent 正在处理上一条消息，请稍候") + return + } + ctx, cancel := context.WithCancelCause(s.ctx) + s.chatBusy[t.ID] = true + s.chatCancel[t.ID] = cancel + s.chatMu.Unlock() + // Persist + broadcast the human turn so the 主 Agent 编排会话 survives page // reloads and updates live: the conversation lives in the activity stream as // worker="mainagent" (the per-task activity table, replayed via SSE). With // attachments, the activity's Detail carries {text, attachments} so the transcript // renders attachment cards. s.engine.emitActivity(t, userActivityWithAttachments("mainagent", req.Message, req.Attachments)) - if ma := s.mainAgentRef(); ma != nil { - // Serialize per task: one main-agent run at a time so concurrent messages - // don't corrupt the shared exp-main transcript. If a prior turn is still - // running, tell the caller to wait rather than starting a racing run. - s.chatMu.Lock() - if s.chatBusy[t.ID] { - s.chatMu.Unlock() - writeErr(w, 409, "主 Agent 正在处理上一条消息，请稍候") - return - } - // Cancellable context so stopChat can abort this turn without affecting - // the server's root context (mirrors the conversation stop pattern). - ctx, cancel := context.WithCancelCause(s.ctx) - s.chatBusy[t.ID] = true - s.chatCancel[t.ID] = cancel - s.chatMu.Unlock() - + var ma *agent.MainAgent + if s.taskRuntimeAvailable(t, "mainagent") { + ma = s.agentsForTask(t).main + } + if ma != nil { // Run the agent on a per-turn cancellable ctx (NOT r.Context()): the turn can // take minutes (multi-tool loop), and binding it to the request lifecycle meant // a page reload / proxy timeout cancelled it mid-run ("context canceled"). The @@ -2461,11 +2989,7 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // handler returns immediately and the browser never needs to hold the request. go func() { defer func() { - cancel(agent.AbortChatTurnFinished) - s.chatMu.Lock() - delete(s.chatBusy, t.ID) - delete(s.chatCancel, t.ID) - s.chatMu.Unlock() + s.finishTaskChat(t.ID, cancel) }() // emit every step (thinking/tool_use/tool_result/text/result) so the main // agent session shows its work live, like worker/planner. The final answer @@ -2477,7 +3001,10 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // taskDir = agent 的工作目录(CWD),与 chatUpload 落盘、ensureRunDir 一致。 taskDir := filepath.Join(s.m.dir, "tasks", t.ID) agentMsg := composeAgentMessage(req.Message, req.Attachments, taskDir) - if _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, agentMsg, emit, t.Notify, resume, t.NotifyGoal); err != nil && ctx.Err() == nil { + s.engine.BeginLLMCall(t.ID) + _, err := ma.Chat(ctx, maTaskID, s.m.Assets(), t.Store, t.Goal, agentMsg, emit, t.Notify, resume, t.NotifyGoal) + s.engine.EndLLMCall(t.ID) + if err != nil && ctx.Err() == nil { s.engine.emitActivity(t, db.Activity{Worker: "mainagent", Kind: "text", IsError: true, Summary: "（主 Agent 出错：" + err.Error() + "）"}) } }() @@ -2486,9 +3013,28 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { } reply := s.fallbackChat(t, req.Message) s.engine.emitActivity(t, db.Activity{Worker: "mainagent", Kind: "text", Summary: reply}) + s.finishTaskChat(t.ID, cancel) writeJSON(w, 200, map[string]any{"reply": reply, "mode": "rule"}) } +func (s *Server) cancelTaskChat(taskID string, cause error) bool { + s.chatMu.Lock() + cancel := s.chatCancel[taskID] + if cancel != nil { + cancel(cause) + } + s.chatMu.Unlock() + return cancel != nil +} + +func (s *Server) finishTaskChat(taskID string, cancel context.CancelCauseFunc) { + cancel(agent.AbortChatTurnFinished) + s.chatMu.Lock() + delete(s.chatBusy, taskID) + delete(s.chatCancel, taskID) + s.chatMu.Unlock() +} + // stopChat aborts the in-flight main-agent turn for a task (manual stop button). // Mirrors pgStopConversation: cancels only the current turn; planner/workers are // unaffected and continue running. The user can send a new message immediately. @@ -2498,14 +3044,10 @@ func (s *Server) stopChat(w http.ResponseWriter, r *http.Request) { writeErr(w, 404, "task not found") return } - s.chatMu.Lock() - cancel := s.chatCancel[t.ID] - s.chatMu.Unlock() - if cancel == nil { + if !s.cancelTaskChat(t.ID, agent.AbortChatStoppedByUser) { writeJSON(w, 200, map[string]any{"status": "idle"}) return } - cancel(agent.AbortChatStoppedByUser) writeJSON(w, 200, map[string]any{"status": "stopping"}) } diff --git a/server/server_mgmt.go b/server/server_mgmt.go index eef78bb..9732b9d 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -5,6 +5,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "log" @@ -66,61 +67,97 @@ func decode(r *http.Request, v any) error { return json.NewDecoder(r.Body).Decod // ---------- tasks (delete) ---------- -// pgDeleteTask removes a task and, per opt-in flags in the (optional) JSON body, -// its associated data: assets (task-owned), findings, recorded traffic (by the -// task's hosts) and the per-task working directory. Empty body → all flags false → -// only the task + its exploration subgraph are removed (legacy behavior). +const taskDeleteDrainTimeout = 10 * time.Second + +func canonicalTaskID(raw string) (string, bool) { + n, err := strconv.ParseInt(raw, 10, 64) + if err != nil || n <= 0 { + return "", false + } + return strconv.FormatInt(n, 10), true +} + func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) { - id := r.PathValue("id") - var opts struct { - Assets bool `json:"assets"` - Findings bool `json:"findings"` - Traffic bool `json:"traffic"` - Files bool `json:"files"` + id, ok := canonicalTaskID(r.PathValue("id")) + if !ok { + writeErr(w, http.StatusBadRequest, "任务 id 无效") + return } - _ = json.NewDecoder(r.Body).Decode(&opts) // optional body; ignore EOF/parse errors - - result := map[string]any{"deleted": id} - nid, _ := strconv.ParseInt(id, 10, 64) - as := s.m.Assets() - - // Traffic first: it needs the task's hosts, which come from assets that a later - // asset-delete would remove. Traffic is host-keyed (no per-task tag), so this - // purges ALL recorded traffic for those hosts, not just this task's. - if opts.Traffic && nid > 0 && as != nil { - if tr := s.m.Traffic(); tr != nil { - if hosts, err := as.HostsByTask(nid); err == nil && len(hosts) > 0 { - if n, err := tr.DeleteHostsExact(hosts); err == nil { - result["traffic_deleted"] = n - } - } + var opts DeleteTaskOptions + if err := decode(r, &opts); err != nil && err != io.EOF { + writeErr(w, 400, "invalid JSON: "+err.Error()) + return + } + if !s.engine.BeginDelete(id) { + writeErr(w, http.StatusConflict, "任务正在删除") + return + } + deleted := false + defer func() { + if !deleted { + s.engine.AbortDelete(id) } - } - if opts.Findings && nid > 0 { - if n, err := s.m.pg.DeleteFindingsByTask(nid); err == nil { - result["findings_deleted"] = n - } - } - if opts.Assets && nid > 0 && as != nil { - if n, err := as.DeleteByTaskID(nid); err == nil { - result["assets_deleted"] = n - } - } + }() - s.engine.Pause(id, agent.AbortTaskDeleted) // cancel running loops before removing - if err := s.m.DeleteTask(id); err != nil { - writeErr(w, 500, err.Error()) + // Main-agent runs use a separate context from planner/workers. Cancel it, then + // wait until both execution domains have returned before removing transcripts. + s.cancelTaskChat(id, agent.AbortTaskDeleted) + drainCtx, cancelDrain := context.WithTimeout(r.Context(), taskDeleteDrainTimeout) + defer cancelDrain() + if err := s.waitTaskQuiescent(drainCtx, id); err != nil { + writeErr(w, http.StatusConflict, "任务仍有运行中的 Agent，删除已取消") return } - // Files last: the per-task working dir /tasks/ (worker artifacts + - // uploads). Confined to the work dir; id is numeric so no traversal risk. - if opts.Files && nid > 0 { - if err := os.RemoveAll(filepath.Join(s.m.dir, "tasks", id)); err == nil { - result["files_deleted"] = true + result, err := s.m.DeleteTask(id, opts) + if err != nil { + var committed *taskDeleteCommittedError + if errors.As(err, &committed) { + // PostgreSQL is already gone. Complete runtime teardown and return the + // auditable counts together with the post-commit cleanup warning. + s.engine.StopTask(id) + s.taskAgentMu.Lock() + delete(s.taskAgents, id) + s.taskAgentMu.Unlock() + deleted = true + writeCommittedTaskDelete(w, result, err) + return + } + writeErr(w, 500, err.Error()) + return + } + // Manager has removed the task from the registry, so no new API operation can + // resolve it. Now stop and join every task-owned Engine goroutine and clear its + // lifecycle maps before releasing the delete barrier. + s.engine.StopTask(id) + s.taskAgentMu.Lock() + delete(s.taskAgents, id) + s.taskAgentMu.Unlock() + deleted = true + writeJSON(w, 200, result) +} + +func writeCommittedTaskDelete(w http.ResponseWriter, result DeleteTaskResult, err error) { + result.CleanupWarning = err.Error() + writeJSON(w, http.StatusOK, result) +} + +func (s *Server) waitTaskQuiescent(ctx context.Context, taskID string) error { + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + s.chatMu.Lock() + chatBusy := s.chatBusy[taskID] + s.chatMu.Unlock() + if s.engine.inflightCount(taskID) == 0 && !chatBusy { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: } } - writeJSON(w, 200, result) } // ---------- agents ---------- @@ -352,6 +389,9 @@ func (s *Server) pgSaveAgentConfig(w http.ResponseWriter, r *http.Request) { if profileChanged { s.invalidateProfileAgents() } + // Task bundles capture operational Agent settings (turn/time budgets, tools, + // web search). Rebuild them even when no global profile is configured. + s.invalidateTaskAgents() // rebuild the live agents so the new max_turns/run_seconds/binding apply without a restart. s.cfgMu.Lock() cfg, on := s.llmCfg, s.llmOn @@ -396,8 +436,8 @@ func (s *Server) pgGetAgent(w http.ResponseWriter, r *http.Request) { } writeJSON(w, 200, map[string]any{ "agent": agentDTO(a), "prompt": cur, "variables": vars, "versions": vers, - "visibility": map[string]any{"mcp": mcp, "skill": sk}, - "llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选 + "visibility": map[string]any{"mcp": mcp, "skill": sk}, + "llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选 "wrapup_prompt": a.WrapupPrompt, // 已保存的收尾提示词(空=用内置默认) "wrapup_default": agent.WrapupDefault(a.Key), // 内置默认(供占位/恢复默认) @@ -1607,7 +1647,9 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) { writeErr(w, 500, err.Error()) return } - // Editing a profile rebuilds any task pinned to it on its next round. + // Editing a profile rebuilds any task pinned to it on its next round. Reapply + // the active profile too: the global fallback and explicit task chains must + // repopulate the same provider-cache entry and therefore share one limiter. s.invalidateProfileAgents() // Hot-apply. Not just when the ACTIVE profile changed: with failover on, any // profile's key/model/priority/pool_exclude edit reshapes the chain, so the @@ -1622,16 +1664,50 @@ func (s *Server) pgDeleteProfile(w http.ResponseWriter, r *http.Request) { return } id, _ := pathInt(r, "id") - if err := pg.DeleteProfile(id); err != nil { - writeErr(w, 500, err.Error()) + if err := pg.DeleteProfileContext(r.Context(), id); err != nil { + switch { + case errors.Is(err, db.ErrActiveLLMProfileDelete): + writeErr(w, 409, "当前激活的 LLM 配置不能删除，请先激活其他配置") + case errors.Is(err, db.ErrLLMProfileReferencesChanged): + writeErr(w, 409, "LLM 配置正在被任务或会话修改，请重试") + case errors.Is(err, context.DeadlineExceeded): + writeErr(w, 409, "等待 LLM 配置引用释放超时，请重试") + case errors.Is(err, db.ErrLLMProfileNotFound): + writeErr(w, 404, err.Error()) + default: + writeErr(w, 500, err.Error()) + } return } s.invalidateProfileAgents() // drop cached agents for the removed profile - s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away) - s.reapplyActiveProfile() // the deleted profile may have been in the failover chain + s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away) + s.restoreTasksAfterProfileDelete(pg) + // Cache invalidation also removed the active profile's shared provider entry. + // Reapply after task-state sync so the wake-up observes the post-delete chain. + s.reapplyActiveProfile() writeJSON(w, 200, map[string]any{"deleted": id}) } +func (s *Server) restoreTasksAfterProfileDelete(pg *db.DB) { + for _, task := range s.m.List() { + if taskID, err := strconv.ParseInt(task.ID, 10, 64); err == nil { + if pt, err := pg.GetTask(taskID); err == nil && pt != nil { + s.syncTaskLLMState(pt) + // Deleting the active entry may advance the task to a ready successor, + // or clear the explicit chain and restore its Agent/global fallback. + // Resume only intents that were blocked by the exhausted chain; a + // paused task remains paused and terminal tasks remain immutable. + if !isTerminalStatus(task.Status) && s.taskRuntimeAvailable(task, "planner", "worker") { + if _, reopenErr := task.Store.ReopenIntentsByBlockedReason(db.IntentBlockedLLMQuota); reopenErr != nil { + log.Printf("[llm-profile] task %s reopen quota-blocked intents after profile delete: %v", task.ID, reopenErr) + } + } + task.Notify() + } + } + } +} + func (s *Server) pgActivateProfile(w http.ResponseWriter, r *http.Request) { pg := s.pg(w) if pg == nil { diff --git a/server/task_delete_barrier_test.go b/server/task_delete_barrier_test.go new file mode 100644 index 0000000..4b95192 --- /dev/null +++ b/server/task_delete_barrier_test.go @@ -0,0 +1,220 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" +) + +func TestEngineDeleteBarrierRejectsNewTaskOperations(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + + if !e.beginTaskOperation(taskID) { + t.Fatal("first task operation should be admitted") + } + if !e.BeginDelete(taskID) { + t.Fatal("first delete should install the barrier") + } + if !e.IsDeleting(taskID) { + t.Fatal("delete barrier should be visible") + } + if !e.IsPaused(taskID) { + t.Fatal("delete should temporarily pause task execution") + } + if e.beginTaskOperation(taskID) { + t.Fatal("operation started after BeginDelete returned") + } + if got := e.inflightCount(taskID); got != 1 { + t.Fatalf("existing operation must remain drainable, inflight=%d", got) + } + e.decInflight(taskID) + + if e.BeginDelete(taskID) { + t.Fatal("duplicate delete must not acquire the barrier") + } + e.AbortDelete(taskID) + if e.IsPaused(taskID) { + t.Fatal("aborted delete should restore a previously running task") + } + if !e.beginTaskOperation(taskID) { + t.Fatal("aborting deletion should reopen operation admission") + } + e.decInflight(taskID) +} + +func TestAbortDeletePreservesExistingPause(t *testing.T) { + e := NewEngine(nil) + e.Pause("42", agent.AbortPausedByUser) + if !e.BeginDelete("42") { + t.Fatal("delete barrier should be installed") + } + e.AbortDelete("42") + if !e.IsPaused("42") { + t.Fatal("aborted delete must preserve a pre-existing user pause") + } +} + +func TestStopTaskCancelsRuntimeAndClearsLifecycleState(t *testing.T) { + e := NewEngine(nil) + const taskID = "42" + e.deleteMu.RLock() + rt := e.registerTaskRoutines(context.Background(), taskID, 1) + e.started.Store(taskID, true) + e.lastAct.Store(taskID, int64(1)) + e.paused.Store(taskID, true) + e.plannerRound.Store(taskID, 3) + e.coordStarted.Store(taskID, true) + e.deleteMu.RUnlock() + + done := make(chan struct{}) + runTaskRoutine(rt, func(ctx context.Context) { + <-ctx.Done() + close(done) + }) + e.StopTask(taskID) + + select { + case <-done: + default: + t.Fatal("task runtime was not cancelled before StopTask returned") + } + if e.Started(taskID) || e.IsPaused(taskID) || e.IsDeleting(taskID) { + t.Fatalf("task lifecycle maps were not cleared: started=%v paused=%v deleting=%v", + e.Started(taskID), e.IsPaused(taskID), e.IsDeleting(taskID)) + } + if _, ok := e.lastAct.Load(taskID); ok { + t.Fatal("last activity state was retained") + } + if _, ok := e.plannerRound.Load(taskID); ok { + t.Fatal("planner round state was retained") + } + if _, ok := e.coordStarted.Load(taskID); ok { + t.Fatal("deadline coordinator state was retained") + } + e.runtimeMu.Lock() + _, retained := e.runtimes[taskID] + e.runtimeMu.Unlock() + if retained { + t.Fatal("task runtime retained after StopTask") + } +} + +func TestCanonicalTaskID(t *testing.T) { + for raw, want := range map[string]string{"42": "42", "00042": "42", "+42": "42"} { + if got, ok := canonicalTaskID(raw); !ok || got != want { + t.Fatalf("canonicalTaskID(%q)=(%q,%v), want (%q,true)", raw, got, ok, want) + } + } + for _, raw := range []string{"", "0", "-1", "abc", "42x"} { + if got, ok := canonicalTaskID(raw); ok { + t.Fatalf("canonicalTaskID(%q)=(%q,true), want invalid", raw, got) + } + } +} + +func TestCancelChatThenWaitTaskQuiescent(t *testing.T) { + e := NewEngine(nil) + chatCtx, chatCancel := context.WithCancelCause(context.Background()) + s := &Server{ + engine: e, + chatBusy: map[string]bool{"7": true}, + chatCancel: map[string]context.CancelCauseFunc{"7": chatCancel}, + } + if !e.beginTaskOperation("7") { + t.Fatal("task operation should be admitted") + } + + cancelObserved := make(chan struct{}) + releaseChat := make(chan struct{}) + go func() { + <-chatCtx.Done() + close(cancelObserved) + <-releaseChat + s.finishTaskChat("7", chatCancel) + }() + if !s.cancelTaskChat("7", agent.AbortTaskDeleted) { + t.Fatal("expected active chat to be cancelled") + } + <-cancelObserved + + waitDone := make(chan error, 1) + waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second) + defer waitCancel() + go func() { waitDone <- s.waitTaskQuiescent(waitCtx, "7") }() + + select { + case err := <-waitDone: + t.Fatalf("wait returned while chat and engine operation were active: %v", err) + case <-time.After(20 * time.Millisecond): + } + close(releaseChat) + select { + case err := <-waitDone: + t.Fatalf("wait returned while engine operation was active: %v", err) + case <-time.After(20 * time.Millisecond): + } + e.decInflight("7") + if err := <-waitDone; err != nil { + t.Fatalf("wait for quiescence: %v", err) + } +} + +func TestWaitTaskQuiescentHonorsContext(t *testing.T) { + s := &Server{ + engine: NewEngine(nil), + chatBusy: map[string]bool{"9": true}, + chatCancel: map[string]context.CancelCauseFunc{}, + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := s.waitTaskQuiescent(ctx, "9"); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected deadline exceeded, got %v", err) + } +} + +func TestTaskChatUploadRejectsDeleteBarrierBeforeCreatingDirectory(t *testing.T) { + dataDir := t.TempDir() + m := &Manager{dir: dataDir, tasks: map[string]*Task{"7": {ID: "7"}}} + s := &Server{m: m, engine: NewEngine(m)} + if !s.engine.BeginDelete("7") { + t.Fatal("delete barrier should be installed") + } + + req := httptest.NewRequest(http.MethodPost, "/api/chat/upload?scope=task&id=7", nil) + rec := httptest.NewRecorder() + s.chatUpload(rec, req) + if rec.Code != http.StatusConflict { + t.Fatalf("expected 409 while deleting, got %d: %s", rec.Code, rec.Body.String()) + } + if _, err := os.Stat(filepath.Join(dataDir, "tasks", "7")); !os.IsNotExist(err) { + t.Fatalf("upload must not recreate task directory, stat err=%v", err) + } +} + +func TestCommittedTaskDeleteReturnsSuccessWithCleanupWarning(t *testing.T) { + rec := httptest.NewRecorder() + result := DeleteTaskResult{Deleted: "7", FilesDeleted: true} + err := &taskDeleteCommittedError{err: errors.New("remove staged task files: permission denied")} + + writeCommittedTaskDelete(rec, result, err) + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String()) + } + var got DeleteTaskResult + if decodeErr := json.NewDecoder(rec.Body).Decode(&got); decodeErr != nil { + t.Fatal(decodeErr) + } + if got.Deleted != "7" || !got.FilesDeleted || !strings.Contains(got.CleanupWarning, "cleanup incomplete") { + t.Fatalf("unexpected response: %+v", got) + } +} diff --git a/server/task_llm.go b/server/task_llm.go new file mode 100644 index 0000000..c1dbe88 --- /dev/null +++ b/server/task_llm.go @@ -0,0 +1,462 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "log" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/llmrec" + "github.com/Autumn-27/norma/llm" + "github.com/Autumn-27/norma/memory" + "github.com/Autumn-27/norma/transcript" +) + +type taskAgentBundle struct { + runtime *taskLLMRuntime // goal decomposition runtime + plannerRuntime *taskLLMRuntime + workerRuntime *taskLLMRuntime + mainRuntime *taskLLMRuntime + pl *agent.Planner + wk *agent.Worker + main *agent.MainAgent +} + +type llmAuditProfile struct { + ID int64 `json:"id"` + Name string `json:"name"` + Format string `json:"format"` + Model string `json:"model"` +} + +type llmTransitionAudit struct { + Mode string `json:"mode"` // automatic | manual | exhausted + Reason string `json:"reason"` + Previous *llmAuditProfile `json:"previous,omitempty"` + Next *llmAuditProfile `json:"next,omitempty"` +} + +type llmActivityMetadata struct { + LLMTransition llmTransitionAudit `json:"llm_transition"` +} + +type taskLLMRuntime struct { + s *Server + taskID string + agentKey string +} + +type taskLLMError struct { + taskID string + chainExhausted bool + cause error +} + +func (e *taskLLMError) Error() string { + if e.chainExhausted { + return fmt.Sprintf("task %s LLM profile chain exhausted: %v", e.taskID, e.cause) + } + return e.cause.Error() +} + +func (e *taskLLMError) Unwrap() error { return e.cause } + +func isTaskLLMRuntimeError(err error) bool { + var target *taskLLMError + return errors.As(err, &target) +} + +func isTaskLLMChainExhausted(err error) bool { + var target *taskLLMError + return errors.As(err, &target) && target.chainExhausted +} + +// isQuotaExhaustedError is intentionally strict. A generic 429, auth error, +// network failure, or 5xx does not rotate providers; the response must explicitly +// identify quota, credits, billing balance, or payment exhaustion. +func isQuotaExhaustedError(err error) bool { + return err != nil && agent.IsQuotaExhaustedMessage(err.Error()) +} + +type taskLLMSelection struct { + task *Task + profileID int64 + revision int64 + provider llm.Provider +} + +type taskLLMStreamHooks struct { + current func() (taskLLMSelection, error) + exhaust func(taskLLMSelection, error) (db.TaskLLMTransition, error) + transition func(taskLLMSelection, db.TaskLLMTransition, error) +} + +func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) { + taskNum, err := parseTaskID(r.taskID) + if err != nil { + return nil, 0, 0, nil, err + } + pt, err := r.s.m.pg.GetTask(taskNum) + if err != nil { + return nil, 0, 0, nil, err + } + if pt == nil { + return nil, 0, 0, nil, fmt.Errorf("task %s not found", r.taskID) + } + r.s.syncTaskLLMState(pt) + t, _ := r.s.m.Task(r.taskID) + if len(pt.LLMProfileIDs) == 0 { + prov, _, ok := r.s.fallbackProviderForAgent(r.agentKey) + if !ok { + return t, 0, pt.LLMChainRevision, nil, fmt.Errorf("task %s has no available fallback LLM provider", r.taskID) + } + return t, 0, pt.LLMChainRevision, prov, nil + } + if pt.ActiveLLMProfileID == nil { + return t, 0, pt.LLMChainRevision, nil, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")} + } + prov, _, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID) + if !ok { + return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, nil, fmt.Errorf("LLM profile #%d is missing or invalid", *pt.ActiveLLMProfileID) + } + return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, prov, nil +} + +func parseTaskID(id string) (int64, error) { + n, err := strconv.ParseInt(id, 10, 64) + if err != nil || n <= 0 { + return 0, fmt.Errorf("invalid task id %q", id) + } + return n, nil +} + +func (r *taskLLMRuntime) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + ctx = llmrec.WithTaskID(ctx, r.taskID) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + task, profileID, revision, provider, err := r.current() + return taskLLMSelection{task: task, profileID: profileID, revision: revision, provider: provider}, err + }, + exhaust: func(selection taskLLMSelection, cause error) (db.TaskLLMTransition, error) { + taskNum, _ := parseTaskID(r.taskID) + transition, err := r.s.m.pg.MarkTaskLLMProfileQuotaExhaustedAtRevision(taskNum, selection.profileID, selection.revision, cause.Error()) + if err != nil { + return transition, err + } + if pt, getErr := r.s.m.pg.GetTask(taskNum); getErr == nil && pt != nil { + r.s.syncTaskLLMState(pt) + } + return transition, nil + }, + transition: func(selection taskLLMSelection, transition db.TaskLLMTransition, cause error) { + r.s.emitTaskLLMTransition(selection.task, transition, cause) + }, + } + return streamTaskLLM(ctx, r.taskID, req, hooks) +} + +func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) iter.Seq2[llm.StreamEvent, error] { + return func(yield func(llm.StreamEvent, error) bool) { + for { + selection, err := hooks.current() + if err != nil { + yield(llm.StreamEvent{}, err) + return + } + committed := false + var pending []llm.StreamEvent + var streamErr error + for event, err := range selection.provider.Stream(ctx, req) { + if err != nil { + streamErr = err + break + } + if !committed && !streamEventCommitsOutput(event) { + pending = append(pending, event) + continue + } + if !committed { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + pending = nil + committed = true + } + if !yield(event, nil) { + return + } + } + if streamErr == nil { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + return + } + // profileID=0 means the explicit chain was cleared while this stable task + // bundle was still in use. Agent/global fallback errors follow the legacy + // behavior and never mutate task failover state. + if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) { + for _, buffered := range pending { + if !yield(buffered, nil) { + return + } + } + yield(llm.StreamEvent{}, streamErr) + return + } + transition, markErr := hooks.exhaust(selection, streamErr) + if markErr != nil { + cause := fmt.Errorf("mark profile quota exhausted after %v: %w", streamErr, markErr) + if committed { + // Output may already have driven tool execution. Report the persistence + // failure, but classify it as router-handled so the worker does not + // replay the entire intent and duplicate those side effects. + yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, cause: cause}) + } else { + yield(llm.StreamEvent{}, cause) + } + return + } + if transition.Advanced && !transition.Stale && hooks.transition != nil { + hooks.transition(selection, transition, streamErr) + } + if committed || (!transition.Stale && transition.NextProfileID == nil) { + yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: streamErr}) + return + } + // No event reached the caller, so replaying the same logical request on + // the next profile cannot duplicate model output or tool execution. + } + } +} + +func streamEventCommitsOutput(event llm.StreamEvent) bool { + switch event.Type { + case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON: + return event.Text != "" + case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop: + return true + default: + return false + } +} + +func (r *taskLLMRuntime) CompactionWindow() int { + taskNum, err := parseTaskID(r.taskID) + if err != nil { + return (agent.Config{}).CompactionWindow() + } + chain, err := r.s.m.pg.TaskLLMProfiles(taskNum) + if err != nil { + return (agent.Config{}).CompactionWindow() + } + if len(chain) == 0 { + if _, cfg, ok := r.s.fallbackProviderForAgent(r.agentKey); ok { + return cfg.CompactionWindow() + } + return (agent.Config{}).CompactionWindow() + } + minimum := 0 + for _, entry := range chain { + if cfg, ok := r.s.loadProfileConfig(entry.ProfileID); ok { + window := cfg.CompactionWindow() + if minimum == 0 || window < minimum { + minimum = window + } + } + } + if minimum == 0 { + return (agent.Config{}).CompactionWindow() + } + return minimum +} + +// fallbackProviderForAgent resolves the pre-existing precedence used by tasks +// without an explicit chain: the Agent's saved binding first, then the current +// global provider (which may come from the default profile or environment). +func (s *Server) fallbackProviderForAgent(agentKey string) (llm.Provider, agent.Config, bool) { + if id := s.effectiveProfileForAgent(agentKey, nil); id != nil { + if prov, cfg, ok := s.providerForProfile(*id); ok { + return s.poolForBinding(*id, prov, cfg), cfg, true + } + } + s.cfgMu.Lock() + defer s.cfgMu.Unlock() + if !s.llmOn || s.llmProv == nil { + return nil, agent.Config{}, false + } + return s.llmProv, s.llmCfg, true +} + +func (s *Server) taskRuntimeAvailable(t *Task, agentKeys ...string) bool { + if t == nil { + return false + } + state := t.llmStateSnapshot() + if len(state.ProfileIDs) > 0 { + if state.ActiveID == nil { + return false + } + _, _, ok := s.providerForProfile(*state.ActiveID) + return ok + } + s.cfgMu.Lock() + globallyReady := s.llmOn && s.llmProv != nil + s.cfgMu.Unlock() + if globallyReady { + return true + } + for _, key := range agentKeys { + if _, _, ok := s.fallbackProviderForAgent(key); !ok { + return false + } + } + return len(agentKeys) > 0 +} + +func (s *Server) invalidateTaskAgents() { + s.taskAgentMu.Lock() + s.taskAgents = map[string]*taskAgentBundle{} + s.taskAgentMu.Unlock() +} + +func (s *Server) agentsForTask(t *Task) *taskAgentBundle { + s.taskAgentMu.Lock() + defer s.taskAgentMu.Unlock() + if bundle := s.taskAgents[t.ID]; bundle != nil { + return bundle + } + goalRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "goals"} + plannerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "planner"} + workerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "worker"} + mainRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "mainagent"} + tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) + window := workerRuntime.CompactionWindow() + wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker")) + wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow) + wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) + wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory"))) + wk.SetWebSearch(s.webSearchFor("worker")) + pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner")) + pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow) + pl.SetKillWork(s.engine.KillWork) + pl.SetSteerWork(s.engine.SteerWork) + pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + pl.SetWebSearch(s.webSearchFor("planner")) + main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent")) + main.SetCompactionWindowResolver(mainRuntime.CompactionWindow) + main.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + main.SetWebSearch(s.webSearchFor("mainagent")) + bundle := &taskAgentBundle{ + runtime: goalRuntime, plannerRuntime: plannerRuntime, workerRuntime: workerRuntime, + mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, + } + s.taskAgents[t.ID] = bundle + return bundle +} + +func (s *Server) syncTaskLLMState(pt *db.Task) { + if pt == nil { + return + } + id := fmt.Sprintf("%d", pt.ID) + s.m.mu.Lock() + if task := s.m.tasks[id]; task != nil { + task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason) + } + s.m.mu.Unlock() +} + +func (s *Server) emitTaskLLMTransition(t *Task, transition db.TaskLLMTransition, cause error) { + if t == nil { + return + } + previous := s.llmAuditProfile(transition.PreviousProfileID) + var next *llmAuditProfile + if transition.NextProfileID != nil { + next = s.llmAuditProfile(*transition.NextProfileID) + } + mode := "automatic" + kind := "llm_switch" + summary := fmt.Sprintf("%s 额度不足", llmAuditProfileLabel(previous)) + if transition.NextProfileID != nil { + summary += fmt.Sprintf("，后续调用切换到 %s", llmAuditProfileLabel(next)) + } else { + mode = "exhausted" + kind = "llm_failover" + summary += "，配置链已耗尽" + } + metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{ + Mode: mode, Reason: cause.Error(), Previous: previous, Next: next, + }}) + s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: kind, IsError: transition.ChainExhausted, Summary: summary, Detail: cause.Error(), Metadata: metadata}) + log.Printf("[llm-failover] task %s: %s", t.ID, summary) +} + +func (s *Server) llmAuditProfile(id int64) *llmAuditProfile { + if id <= 0 || s.m == nil || s.m.pg == nil { + return nil + } + p, err := s.m.pg.ProfileByID(id) + if err != nil || p == nil { + return &llmAuditProfile{ID: id, Name: fmt.Sprintf("配置 #%d", id)} + } + return &llmAuditProfile{ID: p.ID, Name: p.Name, Format: p.Format, Model: p.Model} +} + +func llmAuditProfileLabel(profile *llmAuditProfile) string { + if profile == nil { + return "默认配置" + } + name := profile.Name + if name == "" { + name = fmt.Sprintf("配置 #%d", profile.ID) + } + detail := []string{} + if profile.Format != "" { + detail = append(detail, profile.Format) + } + if profile.Model != "" { + detail = append(detail, profile.Model) + } + if len(detail) == 0 { + return name + } + return fmt.Sprintf("%s（%s）", name, strings.Join(detail, " / ")) +} + +func sameOptionalID(a, b *int64) bool { + if a == nil || b == nil { + return a == nil && b == nil + } + return *a == *b +} + +func (s *Server) emitManualTaskLLMSwitch(t *Task, previousID, nextID *int64) db.Activity { + previous := (*llmAuditProfile)(nil) + next := (*llmAuditProfile)(nil) + if previousID != nil { + previous = s.llmAuditProfile(*previousID) + } + if nextID != nil { + next = s.llmAuditProfile(*nextID) + } + summary := fmt.Sprintf("用户手动将任务 LLM 从 %s 切换到 %s", llmAuditProfileLabel(previous), llmAuditProfileLabel(next)) + metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{ + Mode: "manual", Reason: "用户手动切换任务 LLM", Previous: previous, Next: next, + }}) + return s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "llm_switch", Summary: summary, Detail: summary, Metadata: metadata}) +} diff --git a/server/task_llm_test.go b/server/task_llm_test.go new file mode 100644 index 0000000..faf1cfd --- /dev/null +++ b/server/task_llm_test.go @@ -0,0 +1,507 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "sync" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/norma/harness" + "github.com/Autumn-27/norma/llm" +) + +type scriptedLLMProvider struct { + events []llm.StreamEvent + err error + calls int +} + +func (p *scriptedLLMProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { + return func(yield func(llm.StreamEvent, error) bool) { + p.calls++ + for _, event := range p.events { + if !yield(event, nil) { + return + } + } + if p.err != nil { + yield(llm.StreamEvent{}, p.err) + } + } +} + +func collectTaskLLMStream(seq iter.Seq2[llm.StreamEvent, error]) ([]llm.StreamEvent, error) { + var events []llm.StreamEvent + var streamErr error + for event, err := range seq { + if err != nil { + streamErr = err + break + } + events = append(events, event) + } + return events, streamErr +} + +func TestIsQuotaExhaustedError(t *testing.T) { + t.Parallel() + positive := []string{ + `openai: code=insufficient_quota status=429`, + `provider error: quota_exceeded`, + `grpc code=RESOURCE_EXHAUSTED`, + `You exceeded your current quota, please check your plan and billing details.`, + `billing_hard_limit_reached`, + `billing_not_active`, + `HTTP status 402 Payment Required`, + `账户余额不足，请充值`, + `credit balance is too low`, + } + for _, message := range positive { + if !isQuotaExhaustedError(errors.New(message)) { + t.Errorf("expected quota classification for %q", message) + } + } + + negative := []string{ + `HTTP status 429: rate limit exceeded`, + `RESOURCE_EXHAUSTED: rate limit exceeded`, + `HTTP status 429: quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`, + `RESOURCE_EXHAUSTED: TPM quota exceeded`, + `HTTP status 401: invalid api key`, + `HTTP status 401: insufficient_quota`, + `HTTP status 403: forbidden`, + `HTTP status 403: billing_hard_limit_reached`, + `HTTP status 500: internal server error`, + `HTTP status 503: quota_exceeded`, + `dial tcp: network is unreachable`, + `context length exceeded`, + `invalid request parameter`, + } + for _, message := range negative { + if isQuotaExhaustedError(errors.New(message)) { + t.Errorf("unexpected quota classification for %q", message) + } + } + if isQuotaExhaustedError(nil) { + t.Error("nil error must not be classified as quota exhaustion") + } +} + +func TestTaskDTOUsesEmptyArrays(t *testing.T) { + t.Parallel() + dto := taskDTO(&Task{ID: "7"}, "created") + raw, err := json.Marshal(dto) + if err != nil { + t.Fatal(err) + } + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + t.Fatal(err) + } + for _, field := range []string{"source_task_ids", "llm_profile_ids"} { + value, ok := payload[field].([]any) + if !ok || len(value) != 0 { + t.Fatalf("%s must serialize as [], payload=%s", field, raw) + } + } +} + +func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) { + t.Parallel() + task := &Task{ID: "7"} + var wg sync.WaitGroup + for writer := 0; writer < 2; writer++ { + wg.Add(1) + go func(offset int) { + defer wg.Done() + for i := 0; i < 500; i++ { + id := int64(offset*1000 + i + 1) + task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "") + } + }(writer) + } + for reader := 0; reader < 4; reader++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 500; i++ { + state := task.llmStateSnapshot() + if len(state.ProfileIDs) > 0 { + state.ProfileIDs[0] = -1 + } + _ = taskDTO(task, "running") + } + }() + } + wg.Wait() + state := task.llmStateSnapshot() + if len(state.ProfileIDs) != 2 || state.ProfileIDs[0] <= 0 { + t.Fatalf("snapshot mutation leaked into task state: %+v", state) + } +} + +func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) { + t.Parallel() + quota := errors.New("openai: status 402: insufficient_quota") + first := &scriptedLLMProvider{err: quota} + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "fallback"}}} + active := int64(11) + transitions := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(selection taskLLMSelection, _ error) (db.TaskLLMTransition, error) { + if selection.profileID != 11 { + t.Fatalf("unexpected exhausted profile %d", selection.profileID) + } + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil { + t.Fatal(err) + } + if len(events) != 1 || events[0].Text != "fallback" { + t.Fatalf("unexpected fallback events: %+v", events) + } + if first.calls != 1 || second.calls != 1 || transitions != 1 { + t.Fatalf("calls first=%d second=%d transitions=%d", first.calls, second.calls, transitions) + } +} + +func TestTaskLLMStreamRetriesAfterNonContentStartEvent(t *testing.T) { + t.Parallel() + quota := errors.New("anthropic: status 402: insufficient credit balance") + first := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SEMessageStart}}, + err: quota, + } + second := &scriptedLLMProvider{events: []llm.StreamEvent{ + {Type: llm.SEMessageStart}, + {Type: llm.SETextDelta, Text: "fallback"}, + }} + active := int64(11) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil { + t.Fatal(err) + } + if len(events) != 2 || events[0].Type != llm.SEMessageStart || events[1].Text != "fallback" { + t.Fatalf("non-content event from failed attempt leaked or fallback missing: %+v", events) + } + if first.calls != 1 || second.calls != 1 { + t.Fatalf("calls first=%d second=%d", first.calls, second.calls) + } +} + +func TestTaskLLMStreamRetriesStalePreStreamFailureAgainstReplacementChain(t *testing.T) { + t.Parallel() + first := &scriptedLLMProvider{err: errors.New("status 402: insufficient_quota")} + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "replacement"}}} + current := first + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + profileID := int64(11) + if current == second { + profileID = 22 + } + return taskLLMSelection{profileID: profileID, provider: current}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + current = second + return db.TaskLLMTransition{PreviousProfileID: 11, Stale: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err != nil || len(events) != 1 || events[0].Text != "replacement" { + t.Fatalf("stale call did not continue on replacement chain: events=%+v err=%v", events, err) + } +} + +func TestTaskLLMStreamDoesNotSwitchForOrdinaryErrors(t *testing.T) { + t.Parallel() + tests := []string{ + "openai: status 429: rate limit exceeded", + "openai: status 401: invalid api key", + "openai: status 500: internal server error", + "dial tcp: network is unreachable", + } + for _, message := range tests { + t.Run(message, func(t *testing.T) { + provider := &scriptedLLMProvider{err: errors.New(message)} + exhausted := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + exhausted++ + return db.TaskLLMTransition{}, nil + }, + } + _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if err == nil || err.Error() != message { + t.Fatalf("error=%v, want %q", err, message) + } + if exhausted != 0 { + t.Fatalf("ordinary error changed failover state %d times", exhausted) + } + }) + } +} + +func TestTaskLLMStreamAdvancesAfterMidStreamQuotaWithoutReplay(t *testing.T) { + t.Parallel() + quota := errors.New("anthropic: status 402: credit balance is too low") + first := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "partial"}}, + err: quota, + } + second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "next-call"}}} + active := int64(11) + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + if active == 11 { + return taskLLMSelection{profileID: 11, provider: first}, nil + } + return taskLLMSelection{profileID: 22, provider: second}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + active = 22 + next := active + return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if len(events) != 1 || events[0].Text != "partial" || err == nil { + t.Fatalf("mid-stream result events=%+v err=%v", events, err) + } + if first.calls != 1 || second.calls != 0 { + t.Fatalf("mid-stream call was replayed: first=%d second=%d", first.calls, second.calls) + } + + nextEvents, nextErr := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if nextErr != nil || len(nextEvents) != 1 || nextEvents[0].Text != "next-call" || second.calls != 1 { + t.Fatalf("next call did not use fallback: events=%+v err=%v calls=%d", nextEvents, nextErr, second.calls) + } +} + +func TestTaskLLMStreamMarkFailureAfterToolStartSkipsWorkerReplay(t *testing.T) { + t.Parallel() + quota := errors.New("openai: status 402: insufficient_quota") + markFailure := errors.New("database temporarily unavailable") + provider := &scriptedLLMProvider{ + events: []llm.StreamEvent{{Type: llm.SEToolUseStart}}, + err: quota, + } + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + return db.TaskLLMTransition{}, markFailure + }, + } + + events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if len(events) != 1 || events[0].Type != llm.SEToolUseStart || err == nil { + t.Fatalf("committed tool stream events=%+v err=%v", events, err) + } + if !isTaskLLMRuntimeError(err) { + t.Fatalf("committed mark failure must be router-classified: %T %v", err, err) + } + if retryableWorkerModelError(harness.ReasonModelError, err) { + t.Fatal("worker must not replay after a committed tool stream when quota persistence fails") + } + if provider.calls != 1 { + t.Fatalf("provider calls=%d, want 1", provider.calls) + } +} + +func TestTaskLLMRuntimeErrorSkipsWorkerReplay(t *testing.T) { + partialStreamErr := &taskLLMError{ + taskID: "7", + chainExhausted: false, + cause: errors.New("insufficient_quota after partial output"), + } + if retryableWorkerModelError(harness.ReasonModelError, partialStreamErr) { + t.Fatal("worker must not replay an intent after the task router handled a mid-stream quota error") + } + if !retryableWorkerModelError(harness.ReasonModelError, errors.New("temporary network error")) { + t.Fatal("ordinary model errors should retain the existing worker retry behavior") + } +} + +func TestEngineReadyForTaskOverride(t *testing.T) { + t.Parallel() + e := &Engine{} + if e.Ready() { + t.Fatal("engine without a global provider must not be globally ready") + } + e.SetAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + return &agent.Planner{}, &agent.Worker{} + }) + if !e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("task-specific provider chain must make that task ready") + } +} + +func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) { + t.Parallel() + engine := NewEngine(nil) + engine.started.Store("7", true) + task := &Task{ID: "7", Status: "running"} + profileID := int64(11) + task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted") + s := &Server{engine: engine} + if got := s.resolvedTaskStatus(task); got != "running" { + t.Fatalf("exhausted started status=%q, want running", got) + } + task.Paused = true + if got := s.resolvedTaskStatus(task); got != "paused" { + t.Fatalf("paused exhausted status=%q, want paused", got) + } +} + +func TestAuthoritativeTaskResolverDoesNotFallBackWhenChainUnavailable(t *testing.T) { + t.Parallel() + e := &Engine{} + e.UseLLM(&agent.Planner{}, &agent.Worker{}) + if !e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("global provider should initially make task ready") + } + e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + return nil, nil + }) + if e.ReadyFor(&Task{ID: "7"}) { + t.Fatal("unavailable authoritative task chain must not fall back to global provider") + } +} + +func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) { + t.Parallel() + provider := &scriptedLLMProvider{err: errors.New("billing_not_active")} + transitions := 0 + hooks := taskLLMStreamHooks{ + current: func() (taskLLMSelection, error) { + return taskLLMSelection{profileID: 11, provider: provider}, nil + }, + exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) { + return db.TaskLLMTransition{PreviousProfileID: 11, Advanced: true, ChainExhausted: true}, nil + }, + transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ }, + } + + _, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks)) + if !isTaskLLMChainExhausted(err) { + t.Fatalf("expected chain-exhausted error, got %v", err) + } + if provider.calls != 1 || transitions != 1 { + t.Fatalf("calls=%d transitions=%d", provider.calls, transitions) + } +} + +func TestProfileDeleteRestoresQuotaBlockedIntentWhenFallbackAvailable(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer m.Close() + + profileID, err := m.pg.SaveProfile(&db.LLMProfile{ + Name: fmt.Sprintf("delete-quota-fallback-%d", time.Now().UnixNano()), + Format: "openai", + Model: "test-model", + APIKey: "test-key", + }) + if err != nil { + t.Fatal(err) + } + task, err := m.CreateTaskWithOptions("quota fallback", "resume intent", db.TaskCreateOptions{ + LLMProfileIDs: []int64{profileID}, + }) + if err != nil { + _ = m.pg.DeleteProfile(profileID) + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) + _ = m.pg.DeleteProfile(profileID) + }) + + intentID, err := task.Store.AddIntent(map[string]any{"action": "retry after fallback"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + taskID, err := parseTaskID(task.ID) + if err != nil { + t.Fatal(err) + } + if transition, err := m.pg.MarkTaskLLMProfileQuotaExhausted(taskID, profileID, "insufficient_quota"); err != nil || !transition.ChainExhausted { + t.Fatalf("exhaust task chain: transition=%+v err=%v", transition, err) + } + if err := task.Store.SetIntentBlockedReason(intentID, db.IntentBlockedLLMQuota); err != nil { + t.Fatal(err) + } + task.Paused = true + if err := m.SetTaskPaused(task.ID, true); err != nil { + t.Fatal(err) + } + + // Removing the final explicit entry restores the legacy/global fallback. The + // post-delete sync must reopen only quota-blocked work without unpausing it. + if err := m.pg.DeleteProfile(profileID); err != nil { + t.Fatal(err) + } + s := &Server{ + m: m, + llmOn: true, + llmProv: &scriptedLLMProvider{}, + provByProfile: map[int64]*provEntry{}, + } + s.restoreTasksAfterProfileDelete(m.pg) + + node, err := task.Store.GetNode(intentID) + if err != nil { + t.Fatal(err) + } + if node == nil || node.State != "open" || node.BlockedReason != "" { + t.Fatalf("quota-blocked intent was not reopened: %+v", node) + } + if !task.Paused { + t.Fatal("profile deletion must not resume a paused task") + } + state := task.llmStateSnapshot() + if len(state.ProfileIDs) != 0 || state.FailoverState != "default" { + t.Fatalf("task did not return to fallback state: %+v", state) + } +} diff --git a/traffic/traffic.go b/traffic/traffic.go index ff1d5a3..49675d2 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -10,6 +10,7 @@ import ( "database/sql" "encoding/hex" "encoding/json" + "errors" "fmt" "io/fs" "log" @@ -24,9 +25,9 @@ import ( "sync/atomic" "time" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/permission" actool "github.com/Autumn-27/norma/tool" - "github.com/Autumn-27/artex/db" mproxy "github.com/lqqyt2423/go-mitmproxy/proxy" _ "modernc.org/sqlite" ) @@ -471,28 +472,205 @@ func (t *Traffic) DeleteHost(host string) (int64, error) { // so picking "api.example.com" never sweeps "api.example.com.cn". Returns rows // deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass). func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) { - 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) - if err != nil { - return deleted, err - } - n, _ := res.RowsAffected() - deleted += n - if !removed[h] { - os.RemoveAll(filepath.Join(t.dir, sanitize(h))) - removed[h] = true - } + stage, err := t.StageDeleteHostsExact(hosts) + if err != nil { + return 0, err } - if deleted > 0 { - _ = t.gcBlobs() + deleted := stage.Deleted() + if err := stage.Commit(); err != nil { + return deleted, err } return deleted, nil } +type stagedTrafficPath struct { + source string + staged string +} + +// HostDeleteStage keeps the SQLite delete transaction open and moves matching +// host trees to a same-filesystem staging directory. The Traffic write lock is +// held until Commit or Rollback, so no recorder can recreate a staged tree or a +// blob reference while a task deletion is being coordinated with PostgreSQL. +type HostDeleteStage struct { + traffic *Traffic + tx *sql.Tx + stageDir string + moves []stagedTrafficPath + deleted int64 + done bool +} + +// Deleted reports the number of exchange rows selected by this stage. +func (s *HostDeleteStage) Deleted() int64 { + if s == nil { + return 0 + } + return s.deleted +} + +// StageDeleteHostsExact prepares a reversible exact-host deletion. Callers must +// finish every successful stage with Commit or Rollback. +func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error) { + t.wmu.Lock() + stage := &HostDeleteStage{traffic: t} + fail := func(cause error) (*HostDeleteStage, error) { + if rollbackErr := stage.rollbackLocked(); rollbackErr != nil { + return nil, errors.Join(cause, fmt.Errorf("restore staged traffic: %w", rollbackErr)) + } + return nil, cause + } + + unique := make([]string, 0, len(hosts)) + seen := make(map[string]struct{}, len(hosts)) + for _, host := range hosts { + if _, ok := seen[host]; ok { + continue + } + seen[host] = struct{}{} + unique = append(unique, host) + } + + tx, err := t.db.Begin() + if err != nil { + t.wmu.Unlock() + stage.done = true + return nil, err + } + stage.tx = tx + for _, h := range unique { + res, err := tx.Exec(`DELETE FROM exchanges WHERE host=?`, h) + if err != nil { + return fail(err) + } + n, _ := res.RowsAffected() + stage.deleted += n + } + + stagedSources := make(map[string]struct{}, len(unique)) + for _, host := range unique { + source := filepath.Join(t.dir, sanitize(host)) + if _, duplicatePath := stagedSources[source]; duplicatePath { + continue + } + stagedSources[source] = struct{}{} + if _, err := os.Lstat(source); err != nil { + if os.IsNotExist(err) { + continue + } + return fail(fmt.Errorf("inspect traffic tree for %q: %w", host, err)) + } + if stage.stageDir == "" { + parent := filepath.Join(t.dir, "_delete_staging") + if err := os.MkdirAll(parent, 0o700); err != nil { + return fail(fmt.Errorf("create traffic staging directory: %w", err)) + } + stage.stageDir, err = os.MkdirTemp(parent, "hosts-") + if err != nil { + return fail(fmt.Errorf("create traffic stage: %w", err)) + } + } + staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source))) + if err := os.Rename(source, staged); err != nil { + return fail(fmt.Errorf("stage traffic tree for %q: %w", host, err)) + } + stage.moves = append(stage.moves, stagedTrafficPath{source: source, staged: staged}) + } + return stage, nil +} + +// Rollback restores both the SQLite rows and every staged host tree. +func (s *HostDeleteStage) Rollback() error { + if s == nil || s.done { + return nil + } + return s.rollbackLocked() +} + +func (s *HostDeleteStage) rollbackLocked() error { + if s == nil || s.done { + return nil + } + var errs []error + if s.tx != nil { + if err := s.tx.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) { + errs = append(errs, fmt.Errorf("rollback traffic index: %w", err)) + } + } + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) + } + } + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(errs...) +} + +// Commit makes the staged deletion permanent. SQLite is committed only after +// the caller has committed its PostgreSQL task deletion. +func (s *HostDeleteStage) Commit() error { + if s == nil || s.done { + return nil + } + if err := s.tx.Commit(); err != nil { + // A failed SQLite commit normally rolls the transaction back. Restore the + // trees so the traffic store remains internally consistent and recoverable. + restoreErr := s.restoreTreesLocked() + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(fmt.Errorf("commit traffic index deletion: %w", err), restoreErr) + } + var errs []error + if s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove staged traffic trees: %w", err)) + } + } + if err := s.traffic.gcBlobs(); err != nil { + errs = append(errs, fmt.Errorf("garbage collect traffic blobs: %w", err)) + } + s.done = true + s.traffic.wmu.Unlock() + return errors.Join(errs...) +} + +func (s *HostDeleteStage) restoreTreesLocked() error { + var errs []error + for i := len(s.moves) - 1; i >= 0; i-- { + move := s.moves[i] + if _, err := os.Lstat(move.source); err == nil { + errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source)) + continue + } else if !os.IsNotExist(err) { + errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err)) + continue + } + if err := os.Rename(move.staged, move.source); err != nil { + errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err)) + } + } + if len(errs) == 0 && s.stageDir != "" { + if err := os.RemoveAll(s.stageDir); err != nil { + errs = append(errs, fmt.Errorf("remove traffic stage: %w", err)) + } + } + return errors.Join(errs...) +} + // gcBlobs removes content-addressed blobs in _blobs that no remaining exchange // tree references. References appear as "@blob sha256:" lines inside the // tree's request.http/response.http files (see bodyOrBlob). Best-effort: walk diff --git a/traffic/traffic_test.go b/traffic/traffic_test.go index 25d95b9..497bcf5 100644 --- a/traffic/traffic_test.go +++ b/traffic/traffic_test.go @@ -162,6 +162,152 @@ VALUES(?,?,?,?,?,?,?,?,?,?,?)`, } } +func TestDeleteHostsExactReportsTreeRemovalFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "api.example.com" + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, + "1-0001", 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/1-0001"); err != nil { + t.Fatal(err) + } + + notDir := filepath.Join(dir, "not-a-directory") + if err := os.WriteFile(notDir, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + tr.dir = notDir + + n, err := tr.DeleteHostsExact([]string{host}) + if err == nil { + t.Fatal("DeleteHostsExact returned nil after traffic tree removal failed") + } + if n != 0 { + t.Fatalf("deleted=%d, want 0 after atomic rollback", n) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("rolled-back index count=%d err=%v, want 1", count, err) + } +} + +func TestDeleteHostsExactRollsBackWholeIndexBatch(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + for i, host := range []string{"a.example.com", "b.example.com"} { + id := fmt.Sprintf("1-%04d", i+1) + if err := os.MkdirAll(filepath.Join(dir, host, "GET", id), 0o755); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + } + if _, err := tr.DB().Exec(`CREATE TRIGGER fail_second_host BEFORE DELETE ON exchanges +WHEN OLD.host='b.example.com' BEGIN SELECT RAISE(ABORT, 'forced delete failure'); END`); err != nil { + t.Fatal(err) + } + + if n, err := tr.DeleteHostsExact([]string{"a.example.com", "b.example.com"}); err == nil || n != 0 { + t.Fatalf("DeleteHostsExact failure=(%d,%v), want (0,error)", n, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges`).Scan(&count); err != nil || count != 2 { + t.Fatalf("rolled-back index count=%d err=%v, want 2", count, err) + } + for _, host := range []string{"a.example.com", "b.example.com"} { + if _, err := os.Stat(filepath.Join(dir, host)); err != nil { + t.Fatalf("tree %s changed despite index rollback: %v", host, err) + } + } +} + +func TestStageDeleteHostsExactRollbackRestoresIndexAndTree(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "rollback.example.com" + const id = "1-0001" + tree := filepath.Join(dir, host, "GET", id) + if err := os.MkdirAll(tree, 0o755); err != nil { + t.Fatal(err) + } + marker := filepath.Join(tree, "request.http") + if err := os.WriteFile(marker, []byte("original"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) +VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil { + t.Fatal(err) + } + + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + if stage.Deleted() != 1 { + t.Fatalf("staged deleted=%d, want 1", stage.Deleted()) + } + if _, err := os.Stat(filepath.Join(dir, host)); !os.IsNotExist(err) { + t.Fatalf("host tree was not staged: %v", err) + } + if err := stage.Rollback(); err != nil { + t.Fatal(err) + } + if got, err := os.ReadFile(marker); err != nil || string(got) != "original" { + t.Fatalf("restored tree content=%q err=%v", got, err) + } + var count int + if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 { + t.Fatalf("restored index count=%d err=%v, want 1", count, err) + } +} + +func TestStageDeleteHostsExactRollbackReportsRestoreFailure(t *testing.T) { + dir := t.TempDir() + tr, err := Open(dir, "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tr.Close() + + const host = "conflict.example.com" + hostDir := filepath.Join(dir, host) + if err := os.MkdirAll(hostDir, 0o755); err != nil { + t.Fatal(err) + } + stage, err := tr.StageDeleteHostsExact([]string{host}) + if err != nil { + t.Fatal(err) + } + // Simulate an out-of-band conflicting destination. Rollback must surface the + // failed rename instead of claiming the external data was restored. + if err := os.WriteFile(hostDir, []byte("conflict"), 0o644); err != nil { + t.Fatal(err) + } + if err := stage.Rollback(); err == nil || !strings.Contains(err.Error(), "restore") { + t.Fatalf("rollback err=%v, want restore failure", err) + } + if _, err := os.Stat(stage.stageDir); err != nil { + t.Fatalf("staging was removed after failed restore: %v", err) + } +} + // TestDeleteHostGCBlobs verifies blob garbage collection: after a host's trees // are removed, blobs referenced by no remaining exchange are deleted, while // blobs still referenced (including shared ones) survive. diff --git a/web/src/app/(main)/_components/main-content.tsx b/web/src/app/(main)/_components/main-content.tsx index 3c5a39c..c320a19 100644 --- a/web/src/app/(main)/_components/main-content.tsx +++ b/web/src/app/(main)/_components/main-content.tsx @@ -6,9 +6,9 @@ import { usePathname } from "next/navigation"; import { Separator } from "@/components/ui/separator"; import { SidebarTrigger } from "@/components/ui/sidebar"; +import { useCurrentUser } from "@/hooks/use-current-user"; import { api } from "@/lib/api"; import { cn } from "@/lib/utils"; -import { useCurrentUser } from "@/hooks/use-current-user"; import { AccountSwitcher } from "./sidebar/account-switcher"; import { LayoutControls } from "./sidebar/layout-controls"; @@ -62,7 +62,7 @@ export function MainContent({ children }: { children: ReactNode }) {
{version && ( - 版本 · {version} + 版本 · {version} )} diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index 19f1c2a..d5590e0 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -2,29 +2,20 @@ import * as React from "react"; import Link from "next/link"; -import { useRouter, useSearchParams } from "next/navigation"; +import { useSearchParams } from "next/navigation"; import { toast } from "sonner"; -import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon, Trash2Icon } from "lucide-react"; +import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react"; import { SidebarTrigger } from "@/components/ui/sidebar"; import { Separator } from "@/components/ui/separator"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, - AlertDialogTrigger, -} from "@/components/ui/alert-dialog"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Select, SelectContent, + SelectGroup, SelectItem, SelectTrigger, } from "@/components/ui/select"; @@ -72,9 +63,9 @@ function FieldRow({ } function FindingDetailInner() { - const router = useRouter(); const searchParams = useSearchParams(); const id = searchParams.get("id") ?? ""; + const contextTaskId = searchParams.get("context_task") ?? ""; const [finding, setFinding] = React.useState(null); const [loaded, setLoaded] = React.useState(false); const [tab, setTab] = React.useState("overview"); @@ -85,18 +76,18 @@ function FindingDetailInner() { return; } api - .getFinding(id) + .getFinding(id, contextTaskId || undefined) .then((f) => setFinding(f)) .catch(() => setFinding(null)) .finally(() => setLoaded(true)); - }, [id]); + }, [contextTaskId, id]); React.useEffect(() => { load(); }, [load]); const changeSeverity = React.useCallback( async (next: Severity) => { - if (!finding || next === finding.severity) return; + if (!finding || finding.inherited || next === finding.severity) return; const prev = finding.severity; setFinding({ ...finding, severity: next }); try { @@ -113,7 +104,7 @@ function FindingDetailInner() { const changeStatus = React.useCallback( async (next: FindingStatus) => { - if (!finding || next === finding.status) return; + if (!finding || finding.inherited || next === finding.status) return; const prev = finding.status; setFinding({ ...finding, status: next }); try { @@ -128,18 +119,6 @@ function FindingDetailInner() { [finding, id], ); - // deleteFinding 删除当前漏洞(需二次确认),成功后返回发现列表。 - const deleteFinding = React.useCallback(async () => { - if (!id) return; - try { - await api.deleteFinding(id); - toast.success("已删除漏洞"); - router.push("/function/findings"); - } catch (e) { - toast.error("删除失败：" + (e as Error).message); - } - }, [id, router]); - if (!finding) { return (
@@ -184,25 +163,11 @@ function FindingDetailInner() { - - - - - - - 确认删除该漏洞？ - - 「{title}」将被永久删除，同时从发现列表、任务发现 Tab 与探索图中移除，此操作不可撤销。 - - - - 取消 - 删除 - - - + {finding.inherited && finding.source_task_id && ( + + 来源任务 #{finding.source_task_id} · 只读 + + )}
概览 @@ -271,48 +236,60 @@ function FindingDetailInner() { - {/* 严重等级（可改） */} + {/* 严重等级 */} - changeSeverity(v as Severity)} > - - - - {SEVERITIES.map((sv) => ( - - {statusMeta("severity", sv).label} - - ))} - - + + + + + + {SEVERITIES.map((sv) => ( + + {statusMeta("severity", sv).label} + + ))} + + + + )} - {/* 处理状态（可改） */} + {/* 处理状态 */} - changeStatus(v as FindingStatus)} > - - - - {FINDING_STATUSES.map((st) => ( - - {statusMeta("finding", st).label} - - ))} - - + + + + + + {FINDING_STATUSES.map((st) => ( + + {statusMeta("finding", st).label} + + ))} + + + + )} {/* 漏洞类型 */} diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx index 7d066c6..c890352 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx @@ -152,7 +152,6 @@ const TABS: { key: string; label: string; icon: LucideIcon }[] = [ ]; export function AssetsTab({ taskId }: { taskId: string }) { - // 服务端分页:只取当前分类的当前页,分类总数走 /assets/counts?task_id=,不做数量截断 const [rows, setRows] = React.useState([]); const [total, setTotal] = React.useState(0); const [counts, setCounts] = React.useState>({}); @@ -161,7 +160,6 @@ export function AssetsTab({ taskId }: { taskId: string }) { const [page, setPage] = React.useState(0); const [size, setSize] = React.useState(50); - // 换分类/换页大小:回到第一页,并清掉旧行(各分类列不同,不能串用) React.useEffect(() => { setPage(0); setRows([]); @@ -171,17 +169,16 @@ export function AssetsTab({ taskId }: { taskId: string }) { let alive = true; const load = async () => { try { - const [cur, c] = await Promise.all([ + const [cur, nextCounts] = await Promise.all([ api.taskAssets(taskId, tab, size, page * size), api.assetCounts(taskId), ]); if (!alive) return; setRows(cur.assets); setTotal(cur.total); - setCounts(c ?? {}); + setCounts(nextCounts ?? {}); + setLoaded(true); } catch { - // 保留上一次的数据 - } finally { if (alive) setLoaded(true); } }; @@ -190,7 +187,6 @@ export function AssetsTab({ taskId }: { taskId: string }) { return () => { alive = false; clearInterval(timer); }; }, [taskId, tab, page, size]); - // 轮询期间资产变少时,别停在越界的空页上 React.useEffect(() => { const maxPage = Math.max(0, Math.ceil(total / size) - 1); if (page > maxPage) setPage(maxPage); diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx index 14201fe..e5cfdda 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/coverage-graph-tab.tsx @@ -14,9 +14,9 @@ import { Waypoints, } from "lucide-react"; -import { cn } from "@/lib/utils"; -import { Card, CardContent } from "@/components/ui/card"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { Card, CardContent } from "@/components/ui/card"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Sheet, @@ -26,6 +26,7 @@ import { } from "@/components/ui/sheet"; import { api } from "@/lib/api"; import type { CoverageAssetRef, CoverageAssetRefs, CoverageGraphEdge, CoverageGraphNode } from "@/lib/types"; +import { cn } from "@/lib/utils"; // 每个父节点下、同一类型的子节点默认展示的数量；超出折叠，"展示更多"每次再拉这么多。 const FOLD_LIMIT = 20; @@ -274,10 +275,18 @@ function RefList({ title, items }: { title: string; items: CoverageAssetRef[] })
{items.map((r) => ( -
- #{r.id} - {r.state && {r.state}} - {r.summary || "—"} +
+ #{r.id} + {r.state && {r.state}} + {r.summary || "—"} + {r.inherited && r.source_task_id && ( + + 来源 #{r.source_task_id} · 只读 + + )}
))}
diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx index d4faf24..07d9449 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/findings-tab.tsx @@ -2,25 +2,15 @@ import * as React from "react"; import Link from "next/link"; -import { ArrowUpRightIcon, ChevronRightIcon, Trash2Icon } from "lucide-react"; +import { ArrowUpRightIcon, ChevronRightIcon } from "lucide-react"; import { cn } from "@/lib/utils"; import { Card, CardContent } from "@/components/ui/card"; +import { Badge } from "@/components/ui/badge"; import { StatusBadge } from "@/components/status-badge"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, - AlertDialogTrigger, -} from "@/components/ui/alert-dialog"; -import { Button } from "@/components/ui/button"; import { Select, SelectContent, + SelectGroup, SelectItem, SelectTrigger, } from "@/components/ui/select"; @@ -42,12 +32,12 @@ const FINDING_STATUSES: FindingStatus[] = [ function Row({ f, + contextTaskId, onStatus, - onDelete, }: { f: Finding; + contextTaskId: string; onStatus: (f: Finding, next: FindingStatus) => void; - onDelete: (f: Finding) => void; }) { const [open, setOpen] = React.useState(false); return ( @@ -65,10 +55,17 @@ function Row({ )} /> -
- - {f.name || f.vulnclass || "未分类"} - +
+
+ + {f.name || f.vulnclass || "未分类"} + + {f.inherited && f.source_task_id && ( + + 来源 #{f.source_task_id} · 只读 + + )} +
{f.summary} @@ -92,7 +89,7 @@ function Row({ )}
)} - {f.finding_id ? ( + {f.finding_id && !f.inherited ? ( ) : ( @@ -119,7 +118,11 @@ function Row({ {f.finding_id && ( @@ -130,30 +133,8 @@ function Row({
{open && (
-
- 证据 / PoC - {f.finding_id && ( - - - - - - - 确认删除该漏洞？ - - 「{f.name || f.vulnclass || f.summary || `#${f.finding_id}`}」将被永久删除， - 同时从发现列表、任务发现 Tab 与探索图中移除，此操作不可撤销。 - - - - 取消 - onDelete(f)}>删除 - - - - )} +
+ 证据 / PoC 
             {f.evidence}
@@ -187,7 +168,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
 
   const onStatus = React.useCallback(
     async (f: Finding, next: FindingStatus) => {
-      if (!f.finding_id || next === f.status) return;
+      if (f.inherited || !f.finding_id || next === f.status) return;
       const prev = f.status;
       setFindings((cur) =>
         cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)),
@@ -205,19 +186,8 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     [],
   );
 
-  const onDelete = React.useCallback(async (f: Finding) => {
-    if (!f.finding_id) return;
-    try {
-      await api.deleteFinding(f.finding_id);
-      setFindings((cur) => cur.filter((x) => x.id !== f.id));
-      toast.success("已删除漏洞");
-    } catch (e) {
-      toast.error("删除失败：" + (e as Error).message);
-    }
-  }, []);
-
   const items = findings
-    .filter((f) => f.task_id === taskId)
+    .filter((f) => f.task_id === taskId || f.inherited)
     .sort((a, b) => {
       const order = { critical: 0, high: 1, medium: 2, low: 3 };
       return order[a.severity] - order[b.severity];
@@ -227,11 +197,11 @@ export function FindingsTab({ taskId }: { taskId: string }) {
     
       
         {items.map((f) => (
-          
+          
         ))}
         {items.length === 0 && (
           

-            本任务暂无确认发现。
+            本任务及直接关联任务暂无确认发现。
           

         )}
       
diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx
index 090adb3..056b786 100644
--- a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx
+++ b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx
@@ -1,8 +1,11 @@
 "use client";
 
 import * as React from "react";
-import { FileTextIcon, CopyIcon, CheckIcon } from "lucide-react";
+
+import { CheckIcon, CopyIcon, FileTextIcon } from "lucide-react";
 import { toast } from "sonner";
+
+import { Markdown } from "@/components/markdown";
 import { Button } from "@/components/ui/button";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { api } from "@/lib/api";
@@ -39,6 +42,29 @@ export function ReportTab({ taskId }: { taskId: string }) {
     setTimeout(() => setCopied(false), 1500);
   }
 
+  let content: React.ReactNode;
+  if (loading) {
+    content = (
+      

+        
+        加载中…
+      

+    );
+  } else if (report) {
+    content = (
+      

+        
+      

+    );
+  } else {
+    content = (
+      

+        
+        暂无报告
+      

+    );
+  }
+
   return (
     
       
@@ -53,23 +79,7 @@ export function ReportTab({ taskId }: { taskId: string }) {
           )}
- - {loading ? ( -
- - 加载中… -
- ) : report ? ( - 
-            {report}
-
- ) : ( -
- - 暂无报告 -
- )} - + {content} ); } diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx index 1365089..1e18d1b 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx @@ -1,30 +1,46 @@ "use client"; import * as React from "react"; + import { - RadioIcon, BrainIcon, - UserIcon, - Loader2Icon, - ZapOffIcon, - SendIcon, - SquareIcon, - ClockIcon, CircleCheckIcon, - CircleXIcon, CircleSlashIcon, - ShieldAlertIcon, - WifiOffIcon, - RotateCwIcon, + CircleXIcon, + ClockIcon, + HistoryIcon, + Loader2Icon, PaperclipIcon, + PauseIcon, + PlayIcon, + RadioIcon, + RotateCwIcon, + SendIcon, + ShieldAlertIcon, + SquareIcon, + UserIcon, + WifiOffIcon, XIcon, + ZapOffIcon, + Trash2Icon, } from "lucide-react"; import { toast } from "sonner"; -import { cn } from "@/lib/utils"; -import { Transcript } from "@/components/transcript"; + import { TodoPopover } from "@/components/todo-popover"; +import { Transcript } from "@/components/transcript"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { Textarea } from "@/components/ui/textarea"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupTextarea } from "@/components/ui/input-group"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Tooltip, @@ -43,6 +59,7 @@ import type { TaskNode, TokenTotal, } from "@/lib/types"; +import { cn } from "@/lib/utils"; // fmtBytes renders a human file size for attachment chips (mirrors transcript.tsx). function fmtBytes(n: number): string { @@ -63,6 +80,7 @@ function fmtBytes(n: number): string { // so refresh / tab-switch / sleep / reconnect never drop the newest records. const PAGE = 200; // history page size +const SYSTEM_SCAN_PAGE = 500; // generic activity pages scanned to recover sparse system audit events const MAX_KEEP = 4000; // per-session in-memory cap; older pages re-fetched on scroll-up const STREAM_WINDOW_MS = 5000; // "live" = activity seen within this window @@ -98,6 +116,7 @@ function emptyState(): SessionState { // sessionKeyOf routes an activity to its stable session key. worker="planner" covers // BOTH the Goal Agent's round-0 decomposition and the Planner (single Plan session). function sessionKeyOf(a: Activity): string { + if (a.worker === "system" || a.kind === "llm_switch" || a.kind === "llm_failover") return "system"; if (a.worker === "mainagent") return "main"; if (a.worker === "planner") return "plan"; if (a.intent_id) return `intent:${a.intent_id}`; @@ -122,6 +141,8 @@ function statusIcon(status: SessionStatus) { switch (status) { case "running": // 执行中 return ; + case "paused": + return ; case "pending": // 待领取(open intent) return ; case "done": // 完成 @@ -156,6 +177,7 @@ const roleMeta = { mainagent: { label: "主 Agent", icon: UserIcon }, planner: { label: "规划 Planner", icon: BrainIcon }, worker: { label: "Workers", icon: RadioIcon }, + system: { label: "系统审计", icon: HistoryIcon }, } as const; // The main-agent session is the interactive entry point of this tab and has no @@ -185,10 +207,24 @@ const PLANNER_SESSION: Session = { last_activity: "", }; +// LLM provider switches are task-level audit events rather than agent output. The +// history endpoint has no "system" filter, so this fixed session is populated by +// scanning the generic incremental activity endpoint and then tailed by task SSE. +const SYSTEM_ID = "s-system"; +const SYSTEM_SESSION: Session = { + id: SYSTEM_ID, + role: "system", + title: "系统事件 · LLM 故障转移", + status: "done", + live: false, + last_activity: "", +}; + // keyForSession maps a UI Session → its stable store key (main | plan | intent:). function keyForSession(s: Session): string { if (s.role === "mainagent") return "main"; if (s.role === "planner") return "plan"; + if (s.role === "system") return "system"; return `intent:${s.intent_id}`; } @@ -203,6 +239,8 @@ function intentStatus(state: string): SessionStatus { return "exhausted"; case "stopped": return "stopped"; + case "paused": + return "paused"; case "open": // 待领取，区别于执行中 return "pending"; default: // running @@ -220,9 +258,11 @@ function intentToSession(n: TaskNode): Session { role: "worker", title: label || `Intent ${n.id}`, status: state, - live: state === "running", + live: !n.inherited && state === "running", last_activity: n.ts, intent_id: n.id, + source_task_id: n.source_task_id, + inherited: n.inherited, }; } @@ -233,6 +273,10 @@ function SessionItem({ hasPending, unread, onClick, + onPause, + onResume, + onCancel, + controlling, }: { s: Session; active: boolean; @@ -240,6 +284,10 @@ function SessionItem({ hasPending?: boolean; unread?: number; onClick: () => void; + onPause?: () => void; + onResume?: () => void; + onCancel?: () => void; + controlling?: boolean; }) { const icon = s.role === "worker" ? ( @@ -248,36 +296,57 @@ function SessionItem({ ) : null; + const controllable = s.role === "worker" && !s.inherited && (s.status === "running" || s.status === "paused"); return ( - + {controllable && ( +
+ {s.status === "running" ? ( + + ) : ( + + )} + +
)} - {displayTitle} - {hasPending && ( - - )} - {!active && unread ? ( - - {unread > 99 ? "99+" : unread} - - ) : null} - {s.live && ( - - - 实时 - - )} - +
); } @@ -295,6 +364,8 @@ export function SessionsTab({ taskId }: { taskId: string }) { const [input, setInput] = React.useState(""); const [sending, setSending] = React.useState(false); const [stopping, setStopping] = React.useState(false); + const [controllingIntent, setControllingIntent] = React.useState(null); + const [cancelIntent, setCancelIntent] = React.useState(null); // 方式1 文件上传:选好的附件(已落到任务工作目录 uploads/),随下条消息一起发。 const [attachments, setAttachments] = React.useState([]); const [uploading, setUploading] = React.useState(false); @@ -313,6 +384,42 @@ export function SessionsTab({ taskId }: { taskId: string }) { if (fileInputRef.current) fileInputRef.current.value = ""; } } + + const patchIntentState = React.useCallback((intentId: string, state?: string) => { + const patch = (rows: TaskNode[]) => + state ? rows.map((row) => (row.id === intentId ? { ...row, state } : row)) : rows.filter((row) => row.id !== intentId); + setIntents(patch); + setOlderIntents(patch); + }, []); + + const controlWorker = React.useCallback(async (session: Session, action: "pause" | "resume" | "cancel") => { + if (!session.intent_id || session.inherited || controllingIntent) return; + setControllingIntent(session.intent_id); + try { + const result = await api.controlIntent(taskId, session.intent_id, action); + if (action === "pause") { + patchIntentState(session.intent_id, "paused"); + toast.success(`Worker #${session.intent_id} 已暂停`); + } else if (action === "resume") { + patchIntentState(session.intent_id, "open"); + toast.success(`Worker #${session.intent_id} 已恢复，等待重新领取`); + } else { + patchIntentState(session.intent_id); + if (activeId === session.id) setActiveId(MAIN_ID); + const deleted = result.deleted; + toast.success( + deleted + ? `Worker #${session.intent_id} 已取消，清理 ${deleted.facts} 条事实、${deleted.findings} 个漏洞` + : `Worker #${session.intent_id} 已取消`, + ); + } + } catch (error) { + toast.error("Worker 操作失败：" + (error as Error).message); + } finally { + setControllingIntent(null); + setCancelIntent(null); + } + }, [activeId, controllingIntent, patchIntentState, taskId]); // SSE connection state — surfaced so a dropped realtime link is visible, never // silently shown as "no messages". const [sseLive, setSseLive] = React.useState(false); @@ -326,6 +433,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { const esRef = React.useRef(null); const activeKeyRef = React.useRef("main"); // current session key (for SSE dispatch/unread) const atBottomRef = React.useRef(true); // transcript pinned to bottom? + const llmToastSeqRef = React.useRef>(new Set()); // Per-key request token: a stale response for a key is ignored (guards fast // latest/older interleaving). Writes are ALWAYS keyed, so a late response can only // touch its own session cache — never the currently-viewed one (see §7.5). @@ -353,6 +461,48 @@ export function SessionsTab({ taskId }: { taskId: string }) { const token = (reqTokenRef.current[key] ?? 0) + 1; reqTokenRef.current[key] = token; patchStore(key, (s) => ({ ...s, loading: true, error: undefined })); + + if (key === "system") { + void (async () => { + let since = 0; + let systemItems: Activity[] = []; + for (;;) { + const page = await api.activity(taskId, { since, limit: SYSTEM_SCAN_PAGE }); + if (reqTokenRef.current[key] !== token) return; + systemItems = mergeBySeq( + systemItems, + page.items.filter((item) => item.worker === "system" || item.kind === "llm_failover"), + ); + if (page.items.length < SYSTEM_SCAN_PAGE || page.cursor <= since) break; + since = page.cursor; + } + patchStore(key, (s) => { + const items = mergeBySeq(systemItems, s.items); + return { + ...s, + items, + loaded: true, + loading: false, + hasMore: false, + earliestSeq: items.length ? items[0].seq : 0, + unread: 0, + lastTs: items.length ? items[items.length - 1].ts : s.lastTs, + error: undefined, + }; + }); + })() + .catch((error) => { + if (reqTokenRef.current[key] !== token) return; + patchStore(key, (s) => ({ + ...s, + loading: false, + error: (error as Error).message || "加载失败", + })); + }) + .finally(() => loadingKeysRef.current.delete(key)); + return; + } + api .activityHistory(taskId, key, 0, PAGE) .then((r) => { @@ -424,8 +574,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { .then((r) => { if (alive) setTaskTokens(r.total); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { alive = false; @@ -438,8 +590,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { const load = () => api.interceptTask(taskId) .then((rows) => { if (alive) setPendingIntercepts(rows.filter((r) => r.status === "pending")); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { alive = false; clearInterval(t); }; }, [taskId]); @@ -490,9 +644,15 @@ export function SessionsTab({ taskId }: { taskId: string }) { st.lastTs = st.items.length ? st.items[st.items.length - 1].ts : ""; } buckets.main ??= { ...emptyState(), loaded: true }; + buckets.system ??= { ...emptyState(), loaded: true }; setStore(buckets); }) - .catch(() => setStore({ main: { ...emptyState(), loaded: true } })); + .catch(() => + setStore({ + main: { ...emptyState(), loaded: true }, + system: { ...emptyState(), loaded: true }, + }), + ); return () => { alive = false; }; @@ -532,6 +692,17 @@ export function SessionsTab({ taskId }: { taskId: string }) { } catch { return; // ignore malformed frame } + if ((a.kind === "llm_switch" || a.kind === "llm_failover") && !llmToastSeqRef.current.has(a.seq)) { + llmToastSeqRef.current.add(a.seq); + const transition = a.metadata?.llm_transition; + if (transition?.mode === "exhausted" || a.is_error) { + toast.error(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } else if (transition?.mode === "automatic") { + toast.success(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } else { + toast.info(a.summary, { id: `task-${taskId}-llm-${a.seq}` }); + } + } const k = sessionKeyOf(a); setStore((prev) => { const cur = prev[k] ?? emptyState(); @@ -569,8 +740,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { esRef.current?.close(); esRef.current = null; }; - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [taskId]); + }, [taskId, patchStore]); // ── worker (intent) session list — paged, poll first page lightly ─────────────── React.useEffect(() => { @@ -583,8 +753,10 @@ export function SessionsTab({ taskId }: { taskId: string }) { setIntents(r.items); setIntentsHasMore(r.hasMore); }) - .catch(() => {}); - load(); + .catch(() => { + // Polling is best-effort; the next interval retries automatically. + }); + void load(); const t = setInterval(load, 5000); return () => { active = false; @@ -607,7 +779,9 @@ export function SessionsTab({ taskId }: { taskId: string }) { }); setIntentsHasMore(r.hasMore); }) - .catch(() => {}) + .catch(() => { + // A later manual retry can fetch this page again. + }) .finally(() => setLoadingOlderIntents(false)); }, [taskId, intents, olderIntents, loadingOlderIntents]); @@ -650,6 +824,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { // Worker agent_name format: "work#N · #intentID". Main/planner match by role key. const hasPendingForSession = React.useCallback( (s: Session): boolean => { + if (s.inherited) return false; if (!pendingIntercepts.length) return false; if (s.role === "mainagent") return pendingIntercepts.some((r) => r.agent_name === "mainagent"); if (s.role === "planner") return pendingIntercepts.some((r) => r.agent_name === "planner"); @@ -680,6 +855,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { { ...MAIN_SESSION, live: mainLive }, { ...PLANNER_SESSION, live: plannerLive }, ...workerSessions, + SYSTEM_SESSION, ], [workerSessions, mainLive, plannerLive], ); @@ -688,16 +864,19 @@ export function SessionsTab({ taskId }: { taskId: string }) { mainagent: sessions.filter((s) => s.role === "mainagent"), planner: sessions.filter((s) => s.role === "planner"), worker: sessions.filter((s) => s.role === "worker"), + system: sessions.filter((s) => s.role === "system"), }; const active = sessions.find((s) => s.id === activeId) ?? MAIN_SESSION; const isMain = active.role === "mainagent"; const isPlanner = active.role === "planner"; + const isSystem = active.role === "system"; const activeKey = keyForSession(active); const activeState = store[activeKey]; // Keep the SSE dispatcher's notion of the active session current, and lazily load // + clear unread whenever the active session changes. + // biome-ignore lint/correctness/useExhaustiveDependencies: cache updates must not reactivate the current session. React.useEffect(() => { activeKeyRef.current = activeKey; const st = store[activeKey]; @@ -706,7 +885,6 @@ export function SessionsTab({ taskId }: { taskId: string }) { } else if (st.unread) { patchStore(activeKey, (s) => ({ ...s, unread: 0 })); } - // eslint-disable-next-line react-hooks/exhaustive-deps }, [activeKey]); // Main agent is the human↔orchestrator CONSOLE: only the conversation (user msgs + @@ -721,7 +899,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { const pending = chatExtra.filter((a) => !seen.has(`${a.kind} ${a.summary}`)); return [...items, ...pending].sort((a, b) => a.seq - b.seq); } - if (isPlanner) return items; + if (isPlanner || isSystem) return items; // Worker session: the intent leads the transcript as a right-aligned "user"-style // message (the task handed to this worker), followed by its execution steps. const intentTitle = sessionMeta.get(active.id)?.title ?? active.title; @@ -731,9 +909,23 @@ export function SessionsTab({ taskId }: { taskId: string }) { ts: active.last_activity || "", kind: "intent", // LLM-generated objective — rendered as a distinct (non-human) bubble summary: intentTitle, + source_task_id: active.source_task_id, + inherited: active.inherited, }; return [intentMsg, ...items]; - }, [isMain, isPlanner, activeState, chatExtra, active.id, active.title, active.last_activity, sessionMeta]); + }, [ + isMain, + isPlanner, + isSystem, + activeState, + chatExtra, + active.id, + active.title, + active.last_activity, + active.source_task_id, + active.inherited, + sessionMeta, + ]); // seq of this session's most-recent TodoWrite call — for the Todo popover. const latestTodoSeq = React.useMemo(() => { @@ -791,6 +983,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { null, [], ); + // biome-ignore lint/correctness/useExhaustiveDependencies: activeId intentionally rebinds the scroll listener. React.useEffect(() => { const vp = viewport(); if (!vp) return; @@ -802,6 +995,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { return () => vp.removeEventListener("scroll", onScroll); }, [viewport, activeId, loadEarlier]); // open/switch a session → jump to the latest (bottom) + // biome-ignore lint/correctness/useExhaustiveDependencies: activeId intentionally scrolls a newly selected session. React.useLayoutEffect(() => { const vp = viewport(); if (vp) { @@ -810,6 +1004,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { } }, [activeId, viewport]); // new activity → stick to bottom only if the user is already pinned there + // biome-ignore lint/correctness/useExhaustiveDependencies: activity growth intentionally drives live-edge scrolling. React.useLayoutEffect(() => { if (!atBottomRef.current) return; const vp = viewport(); @@ -819,7 +1014,7 @@ export function SessionsTab({ taskId }: { taskId: string }) { function stop() { if (stopping) return; setStopping(true); - api.stopChat(taskId).finally(() => setStopping(false)); + void api.stopChat(taskId).finally(() => setStopping(false)); } function send() { @@ -906,7 +1101,7 @@ export function SessionsTab({ taskId }: { taskId: string }) {
- {(["mainagent", "planner", "worker"] as const).map((role) => { + {(["mainagent", "planner", "system", "worker"] as const).map((role) => { const items = grouped[role]; if (!items.length) return null; const Meta = roleMeta[role]; @@ -927,11 +1122,16 @@ export function SessionsTab({ taskId }: { taskId: string }) { hasPending={hasPendingForSession(s)} unread={store[keyForSession(s)]?.unread} onClick={() => setActiveId(s.id)} + controlling={controllingIntent === s.intent_id} + onPause={() => void controlWorker(s, "pause")} + onResume={() => void controlWorker(s, "resume")} + onCancel={() => setCancelIntent(s)} /> ); })} {role === "worker" && intentsHasMore && (
)} -
- void pickFiles(e.target.files)} - /> - -