This commit is contained in:
pphua committed 2026-08-24 14:14:15 +08:00
1 parent e67ff26e90
commit 58f23a2773
61 files changed
+11611 -443

No files matched your search

+194
View File
@@ -17,6 +17,200 @@ type ParsedScope struct {
Raw string // original input line
}
<<<<<<< Updated upstream
=======
// ScopeInput is the structured API form for a company scope rule. Empty Kind
// uses the same automatic classification as the single-textarea UI.
type ScopeInput struct {
Kind string `json:"kind,omitempty"`
Value string `json:"value"`
}
// NormalizeICP removes every Unicode whitespace character and folds case. ICP
// matching intentionally performs no fuzzy or punctuation normalization.
func NormalizeICP(value string) string {
return strings.ToLower(strings.Map(func(r rune) rune {
if unicode.IsSpace(r) {
return -1
}
return r
}, strings.TrimSpace(value)))
}
func normalizeKeyword(value string) string {
return strings.ToLower(strings.Join(strings.Fields(value), " "))
}
func looksLikeIPAddress(value string) bool {
value = strings.TrimSpace(value)
if strings.Contains(value, "://") || strings.IndexFunc(value, unicode.IsSpace) >= 0 {
return false
}
if strings.Count(value, ":") >= 2 {
// Require an IPv6-looking prefix. This still catches malformed values such
// as 2001:db8::zz without treating ordinary colon-delimited keywords as IPs.
parts := strings.Split(value, ":")
validSegments := 0
for _, part := range parts {
if part == "" {
if validSegments > 0 || strings.HasPrefix(value, "::") {
return true
}
continue
}
if len(part) > 4 {
return false
}
for _, r := range part {
if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') || (r >= 'A' && r <= 'F')) {
return false
}
}
validSegments++
if validSegments >= 2 {
return true
}
}
return false
}
if !strings.Contains(value, ".") {
return false
}
for _, r := range value {
if r != '.' && (r < '0' || r > '9') {
return false
}
}
return true
}
// ParseScopeInput validates an explicitly typed rule. Legacy callers can omit
// Kind and use the same automatic classification as the single-textarea UI.
func ParseScopeInput(input ScopeInput) (ParsedScope, error) {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
raw := strings.TrimSpace(input.Value)
if kind == "" {
return ParseAutoScopeLine(raw)
}
switch kind {
case "domain", "ip", "cidr":
rule, err := ParseScopeLine(raw)
if err != nil {
return rule, err
}
if rule.Kind != kind {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("%q 不是有效的 %s 范围", raw, kind)
}
return rule, nil
case "icp":
value := NormalizeICP(raw)
if value == "" {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("ICP 不能为空")
}
return ParsedScope{Kind: kind, Value: value, Raw: raw}, nil
case "keyword":
value := normalizeKeyword(raw)
if value == "" {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("企业关键词不能为空")
}
return ParsedScope{Kind: kind, Value: value, Raw: raw}, nil
default:
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("不支持的范围类型: %s", kind)
}
}
// ParseAutoScopeLine classifies one untyped textarea line. Network-looking and
// domain-looking values remain strict so malformed ranges do not silently become
// Agent keywords; all other non-empty text is a keyword.
func ParseAutoScopeLine(line string) (ParsedScope, error) {
raw := strings.TrimSpace(line)
if raw == "" {
return ParsedScope{}, fmt.Errorf("空行")
}
if _, _, err := net.ParseCIDR(raw); err == nil {
return ParseScopeLine(raw)
}
if ip := net.ParseIP(raw); ip != nil {
return ParseScopeLine(raw)
}
if slash := strings.LastIndexByte(raw, '/'); slash > 0 {
address := strings.TrimSpace(raw[:slash])
if net.ParseIP(address) != nil || looksLikeIPAddress(address) {
return ParsedScope{Raw: raw}, fmt.Errorf("无效 CIDR: %s", raw)
}
}
if looksLikeIPAddress(raw) {
return ParsedScope{Raw: raw}, fmt.Errorf("无效 IP: %s", raw)
}
looksLikeDomain := strings.Contains(raw, "://") ||
(strings.Contains(raw, ".") && strings.IndexFunc(raw, unicode.IsSpace) < 0)
if looksLikeDomain {
return ParseScopeLine(raw)
}
// 备案号本身不含点号(如 京ICP备12345678号-1)。带点的文本多半掺了域名或版本号,
// 按 ICP 存下来只会得到一条永远匹配不上任何资产的死规则 —— ICP 归属走的是精确
// 相等比较(见 companies.go 的 kind='icp' 归属查询),所以这类文本归为关键词。
lower := strings.ToLower(raw)
if !strings.ContainsAny(raw, "..。") &&
(strings.Contains(lower, "icp") || strings.Contains(raw, "备案")) {
return ParseScopeInput(ScopeInput{Kind: "icp", Value: raw})
}
return ParseScopeInput(ScopeInput{Kind: "keyword", Value: raw})
}
func scopeHostname(raw string) (string, error) {
candidate := strings.TrimSpace(raw)
if candidate == "" {
return "", fmt.Errorf("主机名为空")
}
if strings.HasPrefix(candidate, "//") {
candidate = "http:" + candidate
} else if !strings.Contains(candidate, "://") {
candidate = "http://" + candidate
}
parsed, err := url.Parse(candidate)
if err != nil || parsed.Host == "" {
if err == nil {
err = fmt.Errorf("缺少主机名")
}
return "", err
}
host := strings.TrimSuffix(strings.TrimSpace(parsed.Hostname()), ".")
if host == "" {
return "", fmt.Errorf("主机名为空")
}
if ip := net.ParseIP(host); ip != nil {
return ip.String(), nil
}
host, err = idna.Lookup.ToASCII(host)
if err != nil {
return "", err
}
host = strings.ToLower(host)
if len(host) > 253 {
return "", fmt.Errorf("域名超过 253 个字符")
}
labels := strings.Split(host, ".")
if len(labels) < 2 {
return "", fmt.Errorf("域名至少需要两个标签")
}
for _, label := range labels {
if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' {
return "", fmt.Errorf("域名标签无效")
}
for _, r := range label {
if !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-') {
return "", fmt.Errorf("域名包含无效字符")
}
}
}
return host, nil
}
>>>>>>> Stashed changes
// ParseScopeLine classifies and validates one scope line (root domain / IP /
// CIDR). Guardrails reject bare TLDs and over-broad networks so a rule can never
// swallow the internet. IP ranges must be expressed as CIDR.
+259
View File
@@ -39,6 +39,265 @@ func TestParseScopeLine(t *testing.T) {
}
}
<<<<<<< Updated upstream
=======
func TestParseAutoScopeLine(t *testing.T) {
cases := []struct {
name string
input string
kind string
normalized string
wantErr bool
}{
{name: "domain", input: "example.com", kind: "domain", normalized: "example.com"},
{name: "url", input: "https://sub.example.com/path", kind: "domain", normalized: "sub.example.com"},
{name: "url query", input: "https://example.com/path?source=x", kind: "domain", normalized: "example.com"},
{name: "url credentials", input: "https://user:[email protected]/path", kind: "domain", normalized: "example.com"},
{name: "url ip", input: "http://203.0.113.10/path", kind: "ip", normalized: "203.0.113.10/32"},
{name: "ipv4", input: "203.0.113.10", kind: "ip", normalized: "203.0.113.10/32"},
{name: "ipv6", input: "2001:db8::10", kind: "ip", normalized: "2001:db8::10/128"},
{name: "cidr", input: "198.51.100.0/24", kind: "cidr", normalized: "198.51.100.0/24"},
{name: "icp latin", input: "京 ICP备 123号", kind: "icp", normalized: "京icp备123号"},
{name: "icp chinese", input: "沪网备案 9988", kind: "icp", normalized: "沪网备案9988"},
{name: "icp domain", input: "icp.example.com", kind: "domain", normalized: "icp.example.com"},
{name: "icp url query", input: "https://example.com/path?icp=1", kind: "domain", normalized: "example.com"},
// 备案号不含点号:掺了域名/版本号的描述性文字归关键词,否则会存成一条
// 永远匹配不上的死 ICP 规则。
{name: "icp with domain text", input: "备案 www.example.com", kind: "keyword", normalized: "备案 www.example.com"},
{name: "icp with version text", input: "某公司 ICP v1.0", kind: "keyword", normalized: "某公司 icp v1.0"},
{name: "icp fullwidth dot", input: "备案 例.com", kind: "keyword", normalized: "备案 例.com"},
{name: "keyword", input: " ACME Security ", kind: "keyword", normalized: "acme security"},
{name: "colon keyword", input: "ACME: Cloud: Security", kind: "keyword", normalized: "acme: cloud: security"},
{name: "empty", input: " ", wantErr: true},
{name: "invalid cidr", input: "10.0.0.0/not-a-prefix", wantErr: true},
{name: "invalid ipv4", input: "999.0.0.1", wantErr: true},
{name: "invalid ipv6", input: "2001:db8::zz", wantErr: true},
{name: "invalid ipv6 cidr", input: "2001:db8::zz/64", wantErr: true},
{name: "overbroad cidr", input: "10.0.0.0/8", wantErr: true},
{name: "bare suffix", input: "co.uk", wantErr: true},
{name: "empty domain label", input: "foo..example.com", wantErr: true},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
rule, err := ParseAutoScopeLine(tc.input)
if tc.wantErr {
if err == nil {
t.Fatalf("ParseAutoScopeLine(%q) = %+v, want error", tc.input, rule)
}
return
}
if err != nil {
t.Fatalf("ParseAutoScopeLine(%q): %v", tc.input, err)
}
if rule.Kind != tc.kind {
t.Fatalf("kind=%q want %q", rule.Kind, tc.kind)
}
got := rule.Value
if rule.Kind == "domain" {
got = rule.Domain
} else if rule.Kind == "ip" || rule.Kind == "cidr" {
got = rule.Net
}
if got != tc.normalized {
t.Fatalf("normalized=%q want %q", got, tc.normalized)
}
})
}
}
func TestExplicitCompanyAttributionSurvivesScopeRebuild(t *testing.T) {
d, as, cs := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
explicitCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Explicit Attribution %d", stamp), "")
if err != nil {
t.Fatal(err)
}
autoCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Automatic Attribution %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, explicitCompany)
defer cleanupCompany(d, autoCompany)
domain := fmt.Sprintf("explicit-%d.invalid", stamp)
icp := fmt.Sprintf("ICP-%d", stamp)
network := fmt.Sprintf("2001:db8:%x::/64", uint64(stamp)&0xffff)
ip := fmt.Sprintf("2001:db8:%x::10", uint64(stamp)&0xffff)
// A pre-existing row with company_id and no provenance value represents old
// installations. The schema default conservatively treats it as explicit.
var assetID int64
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,company_id)
VALUES ('root_domain',$1,$1,$2) RETURNING id`, domain, explicitCompany).Scan(&assetID); err != nil {
t.Fatal(err)
}
defer deleteAsset(d, assetID)
appID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("explicit-app-%d", stamp), ICP: icp, CompanyID: &explicitCompany,
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, appID)
autoAssetID, err := as.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, autoAssetID)
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "icp", Value: icp},
{Kind: "cidr", Value: network},
}
if added, _, invalid, errs := cs.AddScopeInputs(autoCompany, rules, "test"); added != len(rules) || invalid != 0 {
t.Fatalf("AddScopeInputs: added=%d invalid=%d errors=%v", added, invalid, errs)
}
assertCompany := func(id, want int64, source string) {
t.Helper()
var got *int64
var gotSource string
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, id).Scan(&got, &gotSource); err != nil {
t.Fatal(err)
}
if got == nil || *got != want || gotSource != source {
t.Fatalf("asset %d company=%v source=%q, want %d/%q", id, got, gotSource, want, source)
}
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
assertCompany(autoAssetID, autoCompany, "scope")
// Replacing every matching rule detaches only the automatically-owned row.
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("replacement-%d.invalid", stamp)},
}, "test"); invalid != 0 || len(errs) != 0 {
t.Fatalf("UpdateScopeInputs: invalid=%d errors=%v", invalid, errs)
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
var autoCompanyID *int64
var autoSource string
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, autoAssetID).Scan(&autoCompanyID, &autoSource); err != nil {
t.Fatal(err)
}
if autoCompanyID != nil || autoSource != "scope" {
t.Fatalf("automatic asset was not detached: company=%v source=%q", autoCompanyID, autoSource)
}
if err := cs.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
// Deleting the explicitly selected company detaches through the FK, then the
// transactional rebuild may adopt the assets into a still-valid scope.
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "icp", Value: icp},
}, "test"); invalid != 0 || len(errs) != 0 {
t.Fatalf("restore fallback scope: invalid=%d errors=%v", invalid, errs)
}
if err := cs.DeleteCompany(explicitCompany); err != nil {
t.Fatal(err)
}
assertCompany(assetID, autoCompany, "scope")
assertCompany(appID, autoCompany, "scope")
}
func TestParseStructuredCompanyScope(t *testing.T) {
icp, err := ParseScopeInput(ScopeInput{Kind: "ICP", Value: " 京ICP 备 123号-1\t"})
if err != nil {
t.Fatalf("parse ICP: %v", err)
}
if icp.Kind != "icp" || icp.Value != "京icp备123号-1" {
t.Fatalf("unexpected normalized ICP: %+v", icp)
}
keyword, err := ParseScopeInput(ScopeInput{Kind: "keyword", Value: " ACME Security "})
if err != nil {
t.Fatalf("parse keyword: %v", err)
}
if keyword.Value != "acme security" {
t.Fatalf("unexpected normalized keyword: %+v", keyword)
}
if _, err := ParseScopeInput(ScopeInput{Kind: "ip", Value: "example.com"}); err == nil {
t.Fatal("typed IP accepted a domain")
}
}
func TestCompanyICPAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
var suffix int64
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0)+1 FROM companies`).Scan(&suffix); err != nil {
t.Fatal(err)
}
cs := d.Companies()
as := d.Assets()
companyID, _, err := cs.UpsertCompany(fmt.Sprintf("ICP Scope Co %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE task_ids @> ARRAY[$1]::bigint[]`, suffix)
_, _ = d.Exec(`DELETE FROM companies WHERE id=$1`, companyID)
})
// A keyword can guide an Agent, but must never claim an asset by its name.
added, _, invalid, errs := cs.AddScopeInputs(companyID, []ScopeInput{
{Kind: "icp", Value: "京 ICP备 998877号"},
{Kind: "keyword", Value: "ICP Scope"},
}, "unit test")
if added != 2 || invalid != 0 || len(errs) != 0 {
t.Fatalf("add structured scope: added=%d invalid=%d errors=%v", added, invalid, errs)
}
rootID, err := as.UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("icp-scope-%d.example", suffix), ICP: "京icp备998877号", TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
appID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("ICP Scope Keyword Only %d", suffix), TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
icpAppID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("ICP Matched App %d", suffix), ICP: " 京 ICP备 998877号 ", TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
assertCompany := func(assetID int64, want *int64) {
t.Helper()
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, assetID).Scan(&got); err != nil {
t.Fatal(err)
}
if want == nil && got != nil {
t.Fatalf("asset %d attributed by keyword: %d", assetID, *got)
}
if want != nil && (got == nil || *got != *want) {
t.Fatalf("asset %d company=%v want %d", assetID, got, *want)
}
}
assertCompany(rootID, &companyID)
assertCompany(appID, nil)
assertCompany(icpAppID, &companyID)
}
>>>>>>> Stashed changes
// TestCompanyScopeAttribution exercises the full loop against dev PG: create
// company (unique name), add scope, and verify auto-attribution at insert time,
// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
+85
View File
@@ -0,0 +1,85 @@
package db
import (
"fmt"
"time"
)
// Constraint is one operator-authored operation constraint for a task: kind=allow
// (permitted operations) or kind=deny (forbidden operations), free-text. Stored in
// task_constraints, keyed by exploration_id (cascades with the exploration).
type Constraint struct {
ID int64 `json:"id"`
Kind string `json:"kind"` // allow | deny
Text string `json:"text"`
Origin string `json:"origin,omitempty"` // goals | human | system
CreatedAt time.Time `json:"created_at"`
}
// ListConstraints returns this exploration's constraints, allow before deny, oldest
// first within each group (stable render order for the prompt block + UI).
func (s *ExplorationStore) ListConstraints() ([]Constraint, error) {
rows, err := s.db.Query(`
SELECT id, kind, text, COALESCE(origin,''), created_at
FROM task_constraints WHERE exploration_id=$1
ORDER BY (kind='deny'), id`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Constraint
for rows.Next() {
var c Constraint
if err := rows.Scan(&c.ID, &c.Kind, &c.Text, &c.Origin, &c.CreatedAt); err != nil {
return nil, err
}
out = append(out, c)
}
return out, rows.Err()
}
// AddConstraint inserts one constraint (kind must be allow|deny) and returns its id.
func (s *ExplorationStore) AddConstraint(kind, text, origin string) (int64, error) {
if kind != "allow" && kind != "deny" {
return 0, fmt.Errorf("kind 必须是 allow 或 deny")
}
if origin == "" {
origin = "system"
}
var id int64
err := s.db.QueryRow(`
INSERT INTO task_constraints(exploration_id, kind, text, origin)
VALUES ($1, $2, $3, $4) RETURNING id`, s.expID, kind, text, origin).Scan(&id)
return id, err
}
// UpdateConstraint rewrites a constraint's kind + text; scoped to this exploration.
// Returns an error if no such constraint exists.
func (s *ExplorationStore) UpdateConstraint(id int64, kind, text string) error {
if kind != "allow" && kind != "deny" {
return fmt.Errorf("kind 必须是 allow 或 deny")
}
res, err := s.db.Exec(`
UPDATE task_constraints SET kind=$1, text=$2, updated_at=now()
WHERE id=$3 AND exploration_id=$4`, kind, text, id, s.expID)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n == 0 {
return fmt.Errorf("约束不存在")
}
return nil
}
// DeleteConstraint removes a constraint; scoped to this exploration. Returns an
// error if no such constraint exists.
func (s *ExplorationStore) DeleteConstraint(id int64) error {
res, err := s.db.Exec(`DELETE FROM task_constraints WHERE id=$1 AND exploration_id=$2`, id, s.expID)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n == 0 {
return fmt.Errorf("约束不存在")
}
return nil
}
+5 -2
View File
@@ -117,7 +117,8 @@ func intp(v int) *int { return &v }
var builtinAgents = []builtinAgent{
{"goals", "目标拆解", "goals", "把渗透任务目标拆解成若干独立、可验证的子目标。", []promptVar{
{"EngagementDescription", "任务描述(测试对象/背景)", "测试 example.com 站点", "exploration"},
{"Now", "服务端当前时间(RFC3339)", "2026-06-25T00:00:00Z", "runtime"},
// Now 是全局 runtime 变量(见 server.globalPromptVars),不再在各 agent 目录里
// 重复定义,否则 withGlobalVars 追加时会与全局项撞名。
}, false, nil},
{"planner", "规划", "planner", "读取态势、判定目标,只在确有未覆盖的新方向时补充探索意图(每任务一个规划循环)。", []promptVar{
{"Goal", "任务总目标", "拿下 example.com 的管理员权限", "exploration"},
@@ -163,7 +164,9 @@ ON CONFLICT (agent_id, var_name) DO UPDATE
}
// Drop catalog entries for variables that were renamed, so the white-list no
// longer advertises a name templates can't resolve (EngagementTitle→Description).
if _, err := d.Exec(`DELETE FROM agent_prompt_vars WHERE var_name IN ('EngagementTitle', 'CoverageGaps')`); err != nil {
// 'Now' 从各 agent 目录提升为全局 runtime 变量后,旧库里 goals 仍残留一条 'Now'
// 会与全局项撞名(前端变量列表 key 重复);一并清掉。
if _, err := d.Exec(`DELETE FROM agent_prompt_vars WHERE var_name IN ('EngagementTitle', 'CoverageGaps', 'Now')`); err != nil {
return fmt.Errorf("cleanup renamed vars: %w", err)
}
// Default-on interactive_shell for the runtime agents (planner/worker/mainagent/auto)
+335
View File
@@ -205,7 +205,45 @@ func (s *ExplorationStore) SetNodeState(id int64, state string) error {
return err
}
<<<<<<< Updated upstream
// SetIntentState updates an intent node's state (done/blocked/exhausted/open).
=======
// UpdateGoalPayload rewrites a goal node's payload text (and optional vulnclass);
// scoped to kind='goal' so it can never mutate an intent/fact by id. Returns an
// error if no such goal exists in this exploration.
func (s *ExplorationStore) UpdateGoalPayload(id int64, text, vulnclass string) error {
payload := map[string]any{"text": text}
if vulnclass != "" {
payload["vulnclass"] = vulnclass
}
raw, _ := json.Marshal(payload)
res, err := s.db.Exec(`UPDATE exploration_nodes SET payload=$1 WHERE id=$2 AND exploration_id=$3 AND kind='goal'`, string(raw), id, s.expID)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n == 0 {
return fmt.Errorf("目标不存在")
}
return nil
}
// DeleteGoal hard-deletes a goal node; scoped to kind='goal' so it can never drop
// an intent/fact by id. The edges (spawns) and anchors reference it with ON DELETE
// CASCADE, so they go with it; activity.node_id is ON DELETE SET NULL. Returns an
// error if no such goal exists in this exploration.
func (s *ExplorationStore) DeleteGoal(id int64) error {
res, err := s.db.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='goal'`, id, s.expID)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n == 0 {
return fmt.Errorf("目标不存在")
}
return nil
}
// SetIntentState updates an intent node's state.
>>>>>>> Stashed changes
// ResetRunningIntents returns any intent left in 'running' back to 'open' so it
// is re-claimed. Called on startup: a 'running' intent with no live worker (a
// backend restart or crashed worker goroutine left it stuck) would otherwise spin
@@ -258,7 +296,304 @@ WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, termin
return err
}
<<<<<<< Updated upstream
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at`
=======
// CompareAndSetIntentState transitions one local intent only when its current
// state still matches expected. It is the state boundary used by pause/resume and
// worker settlement, so a stale API read cannot overwrite a concurrent finish.
func (s *ExplorationStore) CompareAndSetIntentState(id int64, expected, state string) (bool, error) {
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
res, err := s.db.Exec(`UPDATE exploration_nodes
SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $5 THEN now() ELSE NULL END
WHERE id=$2 AND exploration_id=$3 AND kind='intent' AND state=$4`, state, id, s.expID, expected, terminal)
if err != nil {
return false, err
}
n, err := res.RowsAffected()
return n == 1, err
}
// 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"`
}
// StopIntentWithReason marks a running/paused intent as 'stopped' WITHOUT deleting
// it or any of its yielded facts/findings/activities, and records the user's reason
// as a fact node hanging off that intent (intent --yields--> fact). Returns the new
// fact node id. The caller must first stop a running worker to prevent late writes.
// Used by the "delete work" action, which no longer destroys data — it stops the
// intent and attaches why, so the planner can account for it.
func (s *ExplorationStore) StopIntentWithReason(id int64, reason, origin string) (int64, error) {
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
var state string
var rawPayload []byte
if err := tx.QueryRow(`SELECT state, payload FROM exploration_nodes
WHERE id=$1 AND exploration_id=$2 AND kind='intent' FOR UPDATE`, id, s.expID).Scan(&state, &rawPayload); err != nil {
if err == sql.ErrNoRows {
return 0, fmt.Errorf("intent not found")
}
return 0, err
}
if state != "running" && state != "paused" {
return 0, fmt.Errorf("%w: intent state %s cannot be cancelled", ErrIntentStateConflict, state)
}
// Merge the delete marker into the intent's own payload so it travels with the
// intent everywhere (session list badge, detail banner) without a schema change.
ip := map[string]any{}
_ = json.Unmarshal(rawPayload, &ip)
ip["cancelled_by_user"] = true
ip["cancel_reason"] = reason
newPayload, _ := json.Marshal(ip)
if _, err := tx.Exec(`UPDATE exploration_nodes SET state='stopped', payload=$3, blocked_reason=NULL
WHERE id=$1 AND exploration_id=$2`, id, s.expID, string(newPayload)); err != nil {
return 0, err
}
if origin == "" {
origin = "user"
}
raw, _ := json.Marshal(map[string]any{
"summary": "用户删除了该意图",
"detail": "删除原因:" + reason,
"confidence": "observed",
"user_cancel": true,
})
var factID int64
if err := tx.QueryRow(`
INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, origin)
VALUES ($1, 'fact', $2, 5, 'confirmed', $3) RETURNING id`,
s.expID, string(raw), origin).Scan(&factID); err != nil {
return 0, err
}
if _, err := tx.Exec(`
INSERT INTO exploration_edges(exploration_id, src_id, rel, dst_id) VALUES ($1,$2,$3,$4)
ON CONFLICT DO NOTHING`, s.expID, id, RelYields, factID); err != nil {
return 0, err
}
return factID, tx.Commit()
}
// 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("%w: intent state %s cannot be cancelled", ErrIntentStateConflict, 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')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, 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
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, s.expID, id, RelYields); err != nil {
return out, err
}
// Activity rows are part of the blackboard cleanup, but their token usage is
// irreversible metering data. Fold completed runs plus the latest unfinished
// usage frame into detached daily result rows before deleting the conversation.
// The rows are excluded from per-intent sessions while whole-task totals and
// the original UTC reporting dates remain accurate.
tokenBuckets, err := intentTokenRollup(tx, s.expID, id)
if 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()
for _, bucket := range tokenBuckets {
metadata, _ := json.Marshal(map[string]any{
"cancelled_intent_id": id,
"token_day": bucket.Day.Format(time.DateOnly),
"token_rollup": true,
})
if _, err := tx.Exec(`INSERT INTO activity(
exploration_id, worker, kind, summary, metadata,
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at)
VALUES ($1,'token-ledger','result',$2,$3,$4,$5,$6,$7,$8)`,
s.expID, fmt.Sprintf("已取消意图 #%d 的 Token 计量", id), metadata,
bucket.Usage.InputTokens, bucket.Usage.OutputTokens,
bucket.Usage.CacheReadTokens, bucket.Usage.CacheWriteTokens, bucket.Day); err != nil {
return out, err
}
}
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')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)
)
)`, 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()
}
type tokenUsageBucket struct {
Day time.Time
Usage TokenUsage
}
// intentTokenRollup returns exactly-once usage grouped by its original UTC day.
// Each result terminates one run. A result carrying usage is authoritative; an
// older error/cancel result without usage falls back to the latest cumulative
// usage frame since the previous result. A trailing frame represents a run that
// was interrupted before its terminal result could be persisted.
func intentTokenRollup(tx *sql.Tx, explorationID, intentID int64) ([]tokenUsageBucket, error) {
rows, err := tx.Query(`SELECT kind, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at
FROM activity
WHERE exploration_id=$1 AND node_id=$2 AND kind IN ('usage','result')
ORDER BY id`, explorationID, intentID)
if err != nil {
return nil, err
}
defer rows.Close()
byDay := make(map[time.Time]TokenUsage)
add := func(at time.Time, usage TokenUsage) {
if usage.InputTokens == 0 && usage.OutputTokens == 0 &&
usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
return
}
utc := at.UTC()
day := time.Date(utc.Year(), utc.Month(), utc.Day(), 0, 0, 0, 0, time.UTC)
total := byDay[day]
total.add(usage)
byDay[day] = total
}
var pending TokenUsage
var pendingAt time.Time
hasPending := false
for rows.Next() {
var kind string
var input, output, read, write sql.NullInt64
var createdAt time.Time
if err := rows.Scan(&kind, &input, &output, &read, &write, &createdAt); err != nil {
return nil, err
}
current := TokenUsage{
InputTokens: int(input.Int64),
OutputTokens: int(output.Int64),
CacheReadTokens: int(read.Int64),
CacheWriteTokens: int(write.Int64),
}
hasCurrent := input.Valid || output.Valid || read.Valid || write.Valid
if kind == "usage" {
if hasCurrent {
pending, pendingAt, hasPending = current, createdAt, true
}
continue
}
if hasCurrent {
// Preserve a partially populated legacy terminal row by filling only its
// missing dimensions from the latest cumulative usage frame.
if hasPending {
if !input.Valid {
current.InputTokens = pending.InputTokens
}
if !output.Valid {
current.OutputTokens = pending.OutputTokens
}
if !read.Valid {
current.CacheReadTokens = pending.CacheReadTokens
}
if !write.Valid {
current.CacheWriteTokens = pending.CacheWriteTokens
}
}
add(createdAt, current)
} else if hasPending {
add(createdAt, pending)
}
pending, pendingAt, hasPending = TokenUsage{}, time.Time{}, false
}
if err := rows.Err(); err != nil {
return nil, err
}
if hasPending {
add(pendingAt, pending)
}
buckets := make([]tokenUsageBucket, 0, len(byDay))
for day, usage := range byDay {
buckets = append(buckets, tokenUsageBucket{Day: day, Usage: usage})
}
sort.Slice(buckets, func(i, j int) bool { return buckets[i].Day.Before(buckets[j].Day) })
return buckets, nil
}
func (u *TokenUsage) add(other TokenUsage) {
u.InputTokens += other.InputTokens
u.OutputTokens += other.OutputTokens
u.CacheReadTokens += other.CacheReadTokens
u.CacheWriteTokens += other.CacheWriteTokens
}
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at`
>>>>>>> Stashed changes
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
var n Node
+242 -3
View File
@@ -209,6 +209,244 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
return out, total, err
}
<<<<<<< Updated upstream
=======
// FindingGroup is one task-level bucket in the global findings view. TaskID is
// nil for both findings that never had a task and findings retained after task
// deletion; those records intentionally share one "unassigned/deleted" bucket.
type FindingGroup struct {
TaskID *int64 `json:"task_id"`
TaskName string `json:"task_name"` // 可选任务名称;空=未命名
TaskDescription string `json:"task_description"`
TaskStatus string `json:"task_status"`
Count int `json:"count"`
Critical int `json:"critical"`
High int `json:"high"`
Medium int `json:"medium"`
Low int `json:"low"`
LastFoundAt time.Time `json:"last_found_at"`
}
// ListFindingGroups returns a page of task groups matching the same filters as
// ListFindingsPage. The group count and finding count are independent totals so
// clients can page groups without losing the exact export/selection count.
func (d *DB) ListFindingGroups(f FindingFilter, page, pageSize int) ([]FindingGroup, int, int, error) {
if page <= 0 {
page = 1
}
if pageSize <= 0 {
pageSize = 10
}
where, args := f.where()
grouped := ` FROM findings f LEFT JOIN tasks t ON f.task_id=t.id` + where +
` GROUP BY t.id, t.name, t.description, t.status, t.paused, t.queued`
var groupTotal, findingTotal int
countQuery := `SELECT COUNT(*), COALESCE(SUM(finding_count),0) FROM (` +
`SELECT COUNT(*) AS finding_count` + grouped + `) grouped_findings`
if err := d.QueryRow(countQuery, args...).Scan(&groupTotal, &findingTotal); err != nil {
return nil, 0, 0, err
}
if groupTotal == 0 || page > (groupTotal-1)/pageSize+1 {
return []FindingGroup{}, groupTotal, findingTotal, nil
}
order := "MAX(f.created_at) DESC, t.id DESC NULLS LAST"
if f.Sort == "severity" {
order = `MAX(CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END) DESC, MAX(f.created_at) DESC, t.id DESC NULLS LAST`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
query := fmt.Sprintf(`SELECT t.id, COALESCE(t.name,''), COALESCE(t.description,''), COALESCE(
CASE
WHEN t.status IN ('done','failed','timeout') THEN t.status
WHEN t.queued THEN 'queued'
WHEN t.paused THEN 'paused'
ELSE t.status
END, ''),
COUNT(*),
COUNT(*) FILTER (WHERE f.severity='critical'),
COUNT(*) FILTER (WHERE f.severity='high'),
COUNT(*) FILTER (WHERE f.severity='medium'),
COUNT(*) FILTER (WHERE f.severity='low'),
MAX(f.created_at)%s
ORDER BY %s LIMIT $%d OFFSET $%d`, grouped, order, len(args)+1, len(args)+2)
rows, err := d.Query(query, pageArgs...)
if err != nil {
return nil, 0, 0, err
}
defer rows.Close()
groups := []FindingGroup{}
for rows.Next() {
var group FindingGroup
var taskID sql.NullInt64
if err := rows.Scan(&taskID, &group.TaskName, &group.TaskDescription, &group.TaskStatus, &group.Count,
&group.Critical, &group.High, &group.Medium, &group.Low, &group.LastFoundAt); err != nil {
return nil, 0, 0, err
}
if taskID.Valid {
id := taskID.Int64
group.TaskID = &id
}
groups = append(groups, group)
}
return groups, groupTotal, findingTotal, rows.Err()
}
// ErrFindingOriginUnavailable means a retained finding no longer has a live
// owning task and finding node from which a follow-up intent can be derived.
var ErrFindingOriginUnavailable = errors.New("finding origin is no longer available")
// AddFindingFollowUpIntent atomically creates a priority-10 human intent from a
// live finding node, copies that finding's asset anchors, records the
// finding --derived_from--> intent lineage edge, and persists its audit activity.
// The returned activity is the committed row and can be broadcast as-is without
// calling AppendActivity again.
func (s *ExplorationStore) AddFindingFollowUpIntent(findingID, findingNodeID int64, description string, audit Activity) (int64, Activity, error) {
tx, err := s.db.Begin()
if err != nil {
return 0, Activity{}, err
}
defer tx.Rollback()
var liveNodeID int64
err = tx.QueryRow(`SELECT n.id
FROM findings f
JOIN tasks t ON t.id=f.task_id
JOIN exploration_nodes n ON n.id=f.node_id AND n.exploration_id=t.exploration_id
WHERE f.id=$1 AND f.node_id=$2 AND t.exploration_id=$3 AND n.kind='finding'
FOR SHARE OF f, t, n`, findingID, findingNodeID, s.expID).Scan(&liveNodeID)
if err == sql.ErrNoRows {
return 0, Activity{}, ErrFindingOriginUnavailable
}
if err != nil {
return 0, Activity{}, err
}
anchors := []int64{}
anchorRows, err := tx.Query(`SELECT asset_id FROM exploration_anchors WHERE node_id=$1 ORDER BY asset_id`, liveNodeID)
if err != nil {
return 0, Activity{}, err
}
for anchorRows.Next() {
var assetID int64
if err := anchorRows.Scan(&assetID); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
anchors = append(anchors, assetID)
}
if err := anchorRows.Err(); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
if err := anchorRows.Close(); err != nil {
return 0, Activity{}, err
}
payload := map[string]any{
"summary": description,
"source_finding_id": findingID,
"source_finding_node_id": liveNodeID,
}
if len(anchors) > 0 {
payload["asset_ids"] = anchors
}
raw, err := json.Marshal(payload)
if err != nil {
return 0, Activity{}, err
}
var intentID int64
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
VALUES ($1,'intent',$2,10,'open','human') RETURNING id`, s.expID, raw).Scan(&intentID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id)
SELECT $1, asset_id FROM exploration_anchors WHERE node_id=$2
ON CONFLICT DO NOTHING`, intentID, liveNodeID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id)
VALUES ($1,$2,$3,$4)`, s.expID, liveNodeID, RelDerivedFrom, intentID); err != nil {
return 0, Activity{}, err
}
audit.NodeID = &intentID
if summary := strings.TrimSpace(audit.Summary); summary != "" {
audit.Summary = fmt.Sprintf("%s #%d", summary, intentID)
} else {
audit.Summary = ""
}
metadata := audit.Metadata
if len(metadata) == 0 {
metadata = json.RawMessage(`{}`)
}
if err := tx.QueryRow(`
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, created_at`, s.expID, audit.NodeID, utf8Clean(audit.Worker), utf8Clean(audit.Kind), utf8Clean(audit.Tool), utf8Clean(audit.ToolUseID), audit.IsError,
utf8Clean(audit.Summary), utf8Clean(audit.Detail), metadata, audit.InputTokens, audit.OutputTokens, audit.CacheReadTokens, audit.CacheWriteTokens).
Scan(&audit.ID, &audit.CreatedAt); err != nil {
return 0, Activity{}, err
}
audit.Metadata = metadata
if err := tx.Commit(); err != nil {
return 0, Activity{}, err
}
return intentID, audit, nil
}
// ListFindingsForExport returns findings for the 发现 page 导出功能,携带完整
// report 字段、不分页。ids 非空时按这批 finding id 精确导出(勾选导出),忽略
// filter;ids 为空时按 filter 导出(导出当前筛选/全部)。结果按严重等级降序、
// 再按时间倒序,与「导出汇总报告」的分组顺序一致。
func (d *DB) ListFindingsForExport(f FindingFilter, ids []int64) ([]*DBFinding, error) {
const order = `ORDER BY CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END DESC, f.created_at DESC`
cols := findingSelectCols + `, COALESCE(f.report, '')`
var q string
var args []any
if len(ids) > 0 {
ph := make([]string, len(ids))
for i, id := range ids {
ph[i] = fmt.Sprintf("$%d", i+1)
args = append(args, id)
}
q = `SELECT ` + cols + `
FROM findings f
LEFT JOIN tasks t ON f.task_id = t.id
WHERE f.id IN (` + strings.Join(ph, ",") + `)
` + order
} else {
where, wargs := f.where()
q = `SELECT ` + cols + `
FROM findings f
LEFT JOIN tasks t ON f.task_id = t.id` + where + `
` + order
args = wargs
}
rows, err := d.Query(q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*DBFinding
for rows.Next() {
f := &DBFinding{}
var aidsJSON string
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt,
&f.TaskDescription, &f.Report); err != nil {
return nil, err
}
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
out = append(out, f)
}
return out, rows.Err()
}
>>>>>>> Stashed changes
// FindingStats is the whole-table aggregate powering the 发现 page's stat cards
// and vuln-class filter — computed server-side so it stays exact regardless of
// pagination.
@@ -229,6 +467,7 @@ type FindingStats struct {
// the id.
type FindingTaskOption struct {
ID int64 `json:"id"`
Name string `json:"name"` // 可选任务名称;空=未命名
Description string `json:"description"`
Count int `json:"count"`
}
@@ -265,11 +504,11 @@ func (d *DB) FindingStats() (*FindingStats, error) {
}
// 任务下拉:有漏洞的任务,带描述(任务删除后为空,前端回退 id)和条数,最新有漏洞的排前。
trows, err := d.Query(`SELECT f.task_id, COALESCE(t.description, ''), COUNT(*)
trows, err := d.Query(`SELECT f.task_id, COALESCE(t.name, ''), COALESCE(t.description, ''), COUNT(*)
FROM findings f
LEFT JOIN tasks t ON f.task_id = t.id
WHERE f.task_id IS NOT NULL
GROUP BY f.task_id, t.description
GROUP BY f.task_id, t.name, t.description
ORDER BY MAX(f.created_at) DESC`)
if err != nil {
return nil, err
@@ -277,7 +516,7 @@ func (d *DB) FindingStats() (*FindingStats, error) {
defer trows.Close()
for trows.Next() {
var opt FindingTaskOption
if err := trows.Scan(&opt.ID, &opt.Description, &opt.Count); err != nil {
if err := trows.Scan(&opt.ID, &opt.Name, &opt.Description, &opt.Count); err != nil {
return nil, err
}
st.Tasks = append(st.Tasks, opt)
+23 -12
View File
@@ -34,6 +34,7 @@ type InterceptPending struct {
ToolName string `json:"tool_name"`
ToolInput json.RawMessage `json:"tool_input"`
Status string `json:"status"`
Reason string `json:"reason"` // 规则 message 或模型判定理由(前缀 [模型])
DecidedAt *time.Time `json:"decided_at"`
CreatedAt time.Time `json:"created_at"`
}
@@ -105,7 +106,7 @@ func (d *DB) ToggleInterceptRule(id int64, enabled bool) error {
// CreateInterceptPending inserts a pending approval record and returns its ID.
// convID == 0 → conversation_id stored as NULL (background task).
// taskID == "" → task_id stored as NULL.
func (d *DB) CreateInterceptPending(ruleID, convID int64, taskID, agentName, toolName string, input []byte) (int64, error) {
func (d *DB) CreateInterceptPending(ruleID, convID int64, taskID, agentName, toolName string, input []byte, reason string) (int64, error) {
raw := json.RawMessage(input)
if len(raw) == 0 {
raw = json.RawMessage("{}")
@@ -118,11 +119,16 @@ func (d *DB) CreateInterceptPending(ruleID, convID int64, taskID, agentName, too
if taskID != "" {
taskIDPtr = &taskID
}
// ruleID == 0 → NULL: the LLM fallback judge has no owning rule.
var ruleIDPtr *int64
if ruleID != 0 {
ruleIDPtr = &ruleID
}
var id int64
err := d.QueryRow(`
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input)
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id`,
ruleID, convIDPtr, taskIDPtr, agentName, toolName, raw).Scan(&id)
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, reason)
VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id`,
ruleIDPtr, convIDPtr, taskIDPtr, agentName, toolName, raw, reason).Scan(&id)
return id, err
}
@@ -136,7 +142,7 @@ func (d *DB) DecideInterceptPending(id int64, status string) error {
// (status = 'allowed' | 'denied'), decided_at stamped now. Used to log allow/deny
// rule matches for observability — they don't block and need no user action, so unlike
// CreateInterceptPending (which starts 'pending') this records the outcome directly.
func (d *DB) CreateDecidedIntercept(ruleID, convID int64, taskID, agentName, toolName string, input []byte, status string) (int64, error) {
func (d *DB) CreateDecidedIntercept(ruleID, convID int64, taskID, agentName, toolName string, input []byte, status, reason string) (int64, error) {
raw := json.RawMessage(input)
if len(raw) == 0 {
raw = json.RawMessage("{}")
@@ -149,19 +155,24 @@ func (d *DB) CreateDecidedIntercept(ruleID, convID int64, taskID, agentName, too
if taskID != "" {
taskIDPtr = &taskID
}
// ruleID == 0 → NULL: the LLM fallback judge has no owning rule.
var ruleIDPtr *int64
if ruleID != 0 {
ruleIDPtr = &ruleID
}
var id int64
err := d.QueryRow(`
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, decided_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, NOW()) RETURNING id`,
ruleID, convIDPtr, taskIDPtr, agentName, toolName, raw, status).Scan(&id)
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, reason, decided_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW()) RETURNING id`,
ruleIDPtr, convIDPtr, taskIDPtr, agentName, toolName, raw, status, reason).Scan(&id)
return id, err
}
const interceptPendingCols = `id, rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, decided_at, created_at`
const interceptPendingCols = `id, rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, reason, decided_at, created_at`
func scanInterceptPending(s interface{ Scan(...any) error }, p *InterceptPending) error {
return s.Scan(&p.ID, &p.RuleID, &p.ConversationID, &p.TaskID, &p.AgentName,
&p.ToolName, &p.ToolInput, &p.Status, &p.DecidedAt, &p.CreatedAt)
&p.ToolName, &p.ToolInput, &p.Status, &p.Reason, &p.DecidedAt, &p.CreatedAt)
}
// ListPendingIntercepts returns all unresolved approval requests, newest first.
@@ -206,14 +217,14 @@ type InterceptApprovalRow struct {
func scanInterceptApprovalRow(rows interface{ Scan(...any) error }, r *InterceptApprovalRow) error {
return rows.Scan(
&r.ID, &r.RuleID, &r.ConversationID, &r.TaskID, &r.AgentName,
&r.ToolName, &r.ToolInput, &r.Status, &r.DecidedAt, &r.CreatedAt,
&r.ToolName, &r.ToolInput, &r.Status, &r.Reason, &r.DecidedAt, &r.CreatedAt,
&r.ConvTitle, &r.ConvAgentKey, &r.RuleName,
)
}
const approvalRowSelect = `
SELECT ip.id, ip.rule_id, ip.conversation_id, ip.task_id, ip.agent_name,
ip.tool_name, ip.tool_input, ip.status, ip.decided_at, ip.created_at,
ip.tool_name, ip.tool_input, ip.status, ip.reason, ip.decided_at, ip.created_at,
COALESCE(c.title,'') AS conv_title,
COALESCE(c.agent_key,'') AS conv_agent_key,
COALESCE(ir.name,'') AS rule_name
+104
View File
@@ -168,6 +168,21 @@ CREATE TABLE IF NOT EXISTS exploration_anchors (
);
CREATE INDEX IF NOT EXISTS idx_anchor_asset ON exploration_anchors(asset_id);
-- task_constraints: operator-authored operation constraints (allow/deny) for a task.
-- Extracted by the goals decomposer at round 0 (from goal/description), editable at
-- runtime by the main agent + 总览「约束管理」. Injected into the planner/worker system
-- prompt each round (config-gated) to keep exploration within the operator's boundary.
CREATE TABLE IF NOT EXISTS task_constraints (
id BIGSERIAL PRIMARY KEY,
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('allow','deny')),
text TEXT NOT NULL,
origin TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_task_constraints_exp ON task_constraints(exploration_id);
CREATE TABLE IF NOT EXISTS activity (
id BIGSERIAL PRIMARY KEY,
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
@@ -227,6 +242,7 @@ CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles
-- =====================================================================
CREATE TABLE IF NOT EXISTS tasks (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL DEFAULT '',
description TEXT NOT NULL,
goal TEXT NOT NULL,
exploration_id BIGINT NOT NULL UNIQUE
@@ -239,6 +255,7 @@ CREATE TABLE IF NOT EXISTS tasks (
parent_ref TEXT,
timeout_seconds INTEGER NOT NULL DEFAULT 0,
plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300,
coverage_enabled BOOLEAN NOT NULL DEFAULT true,
first_run_at TIMESTAMPTZ,
deadline_at TIMESTAMPTZ,
deleted_at TIMESTAMPTZ,
@@ -253,6 +270,89 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- planner 心跳触发间隔(秒);补旧库。默认 300s(5min)。见 docs/planner-trigger-impl-plan.md
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300;
<<<<<<< Updated upstream
=======
-- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false;
-- 资产覆盖度功能开关;补旧库。true(默认)=计算/展示测试覆盖度、自动累积测试范围、
-- 给 agent 开放 add_task_scope/list_untested_assets;false=全部关闭(见 task_scope.go)。
-- 存量任务默认 true 保持原行为;company 关联(task_scope kind=company)不受此开关影响。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS coverage_enabled BOOLEAN NOT NULL DEFAULT true;
-- queued_at makes admission FIFO reflect the actual enqueue order rather than the
-- task creation order. queue_mode distinguishes first bootstrap from resuming an
-- exploration that already owns goals/history.
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued_at TIMESTAMPTZ;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queue_mode TEXT NOT NULL DEFAULT '';
-- 可选的任务名称;补旧库。空串=未命名,前端展示时回退到描述。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS name TEXT NOT NULL DEFAULT '';
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;
-- Reusable task description/goal presets. nkey is the normalized, case-insensitive
-- identity used to reject visually equivalent duplicate names.
CREATE TABLE IF NOT EXISTS task_templates (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
nkey TEXT NOT NULL UNIQUE,
description TEXT NOT NULL,
goal TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_task_templates_updated ON task_templates(updated_at DESC, id DESC);
DROP TRIGGER IF EXISTS trg_task_templates_upd ON task_templates;
CREATE TRIGGER trg_task_templates_upd BEFORE UPDATE ON task_templates
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- 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);
>>>>>>> Stashed changes
-- 任务测试范围(资产覆盖度的分母 + 授权边界)。
-- 自动填(source='auto'):insertAssets 顶层按 worker 显式插入的资产类型加保守范围
@@ -526,11 +626,15 @@ CREATE TABLE IF NOT EXISTS intercept_pending (
tool_input JSONB NOT NULL DEFAULT '{}',
status TEXT NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'allowed', 'denied', 'timeout')),
-- 判定理由:规则命中时为规则 message;LLM 兜底判定时为模型给的简短理由(前缀 [模型])。
reason TEXT NOT NULL DEFAULT '',
decided_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_intercept_pending_status ON intercept_pending(status, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_intercept_pending_task ON intercept_pending(task_id, created_at DESC);
-- 补旧库:reason 列(已发版,加列要带 IF NOT EXISTS)。
ALTER TABLE intercept_pending ADD COLUMN IF NOT EXISTS reason TEXT NOT NULL DEFAULT '';
-- =====================================================================
-- L. 漏洞发现持久化
+216
View File
@@ -0,0 +1,216 @@
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(c.name,''), 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
LEFT JOIN companies c ON c.id=ts.company_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.CompanyName, &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
}
+28 -7
View File
@@ -16,10 +16,13 @@ type TaskScope struct {
TaskID int64 `json:"task_id"`
Kind string `json:"kind"` // company|root_domain|subdomain|ip|cidr
CompanyID *int64 `json:"company_id,omitempty"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Source string `json:"source"`
Reason string `json:"reason,omitempty"`
// CompanyName is resolved for kind=company so callers can label a scope row
// without a second lookup. Empty when the row is not a company reference.
CompanyName string `json:"company_name,omitempty"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Source string `json:"source"`
Reason string `json:"reason,omitempty"`
}
// stripHostPort drops a trailing :port from a host:port / ip:port / [ipv6]:port
@@ -187,8 +190,11 @@ func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (Ta
// ListTaskScope returns all scope rows for a task.
func (s *AssetStore) ListTaskScope(taskID int64) ([]TaskScope, error) {
rows, err := s.db.Query(`
SELECT id, kind, COALESCE(company_id,0), COALESCE(domain,''), COALESCE(net::text,''), source, COALESCE(reason,'')
FROM task_scope WHERE task_id=$1 ORDER BY id`, taskID)
SELECT ts.id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''),
COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'')
FROM task_scope ts
LEFT JOIN companies c ON c.id=ts.company_id
WHERE ts.task_id=$1 ORDER BY ts.id`, taskID)
if err != nil {
return nil, err
}
@@ -197,7 +203,7 @@ FROM task_scope WHERE task_id=$1 ORDER BY id`, taskID)
for rows.Next() {
var t TaskScope
var cid int64
if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.Domain, &t.Net, &t.Source, &t.Reason); err != nil {
if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.CompanyName, &t.Domain, &t.Net, &t.Source, &t.Reason); err != nil {
return nil, err
}
t.TaskID = taskID
@@ -227,6 +233,7 @@ type CoverageByType struct {
// NOT a precise metric. Denominator = assets matching any active task_scope row;
// Tested = those anchored to at least one fact node in the exploration.
type Coverage struct {
Enabled bool `json:"enabled"` // 资产覆盖度功能是否开启;false 时其余字段为零值
ScopeRows int `json:"scope_rows"` // 0 → 范围未锚定
Denominator int `json:"denominator"` // 范围内资产数
Tested int `json:"tested"` // 已测(约)
@@ -234,6 +241,20 @@ type Coverage struct {
ByType []CoverageByType `json:"by_type"` // 按资产类型的 总数/已测
}
// CoverageEnabled reports whether a task has the asset-coverage feature turned on
// (tasks.coverage_enabled). Missing row / error → true (fail open to the default),
// so unknown/legacy tasks keep the historical behavior. taskID<=0 → true.
func (s *AssetStore) CoverageEnabled(taskID int64) bool {
if taskID <= 0 {
return true
}
var enabled bool
if err := s.db.QueryRow(`SELECT COALESCE(coverage_enabled,true) FROM tasks WHERE id=$1`, taskID).Scan(&enabled); err != nil {
return true
}
return enabled
}
// task_scope→assets match predicate, reused by the count / by-type / untested queries. $1=taskID.
const covTargetCTE = `
target AS (
+139
View File
@@ -8,6 +8,7 @@ import (
// Task is a row in the task registry (1:1 with an exploration).
type Task struct {
ID int64 `json:"id"`
Name string `json:"name"` // 可选任务名称;空=未命名
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
@@ -24,6 +25,11 @@ type Task struct {
// planner 心跳触发间隔(秒):距上轮 plan 结束/任务开始满该值且期间无触发 → 触发一轮。
// 下限=默认=300(5min),低于一律抬到 300(在 CreateTask 归一)。见 docs/planner-trigger-impl-plan.md
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"`
// CoverageEnabled 是「资产覆盖度功能」总开关(默认 true)。false 时:不计算/不展示测试
// 覆盖度、不自动累积 task_scope(source=auto)、不给 agent 开放 add_task_scope/
// list_untested_assets、态势里不注入 coverage 块(scope 字段仍保留)。company 关联
// (task_scope kind=company)与此开关无关,永不受影响。见 db/task_scope.go。
CoverageEnabled bool `json:"coverage_enabled"`
}
// IsTerminal reports whether a task status is a terminal (finished) state.
@@ -47,6 +53,43 @@ func normalizeHeartbeat(sec int) int {
}
func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) {
<<<<<<< Updated upstream
=======
var ids []int64
if llmProfileID != nil {
ids = []int64{*llmProfileID}
}
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 {
Name string // 可选任务名称;空=未命名
SourceTaskIDs []int64
CompanyIDs []int64
LLMProfileIDs []int64
TimeoutSeconds int
PlanHeartbeatSeconds int
// CoverageEnabled 是「资产覆盖度功能」开关;nil=默认开(true),让不关心该开关的创建
// 路径(编排 spawn、老 API)沿用原行为。仅 web 创建任务时可显式传 false 关闭。
CoverageEnabled *bool
}
// 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)
}
companyIDs, err := NormalizeTaskCompanyIDs(opts.CompanyIDs)
if err != nil {
return nil, err
}
opts.CompanyIDs = companyIDs
>>>>>>> Stashed changes
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -71,24 +114,120 @@ INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, or
VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil {
return nil, err
}
<<<<<<< Updated upstream
if timeoutSeconds < 0 {
timeoutSeconds = 0
=======
if opts.TimeoutSeconds < 0 {
opts.TimeoutSeconds = 0
}
opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds)
coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled
var active *int64
if len(opts.LLMProfileIDs) > 0 {
id := opts.LLMProfileIDs[0]
active = &id
}
t := &Task{
Name: opts.Name,
Description: description, Goal: goal, ExplorationID: expID,
LLMProfileID: active, ActiveLLMProfileID: active,
LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...),
SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), opts.CompanyIDs...),
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
CoverageEnabled: coverageEnabled,
>>>>>>> Stashed changes
}
planHeartbeatSeconds = normalizeHeartbeat(planHeartbeatSeconds) // <=0→300,<30→30
t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds}
if err := tx.QueryRow(`
<<<<<<< Updated upstream
INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds, plan_heartbeat_seconds) VALUES ($1,$2,$3,$4,$5,$6)
RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds, planHeartbeatSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
=======
INSERT INTO tasks(name, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled)
VALUES ($1,$2,$3,$4,$5,$5,$6,$7,$8)
RETURNING id, status, paused, created_at`, opts.Name, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
>>>>>>> Stashed changes
return nil, err
}
return t, tx.Commit()
}
<<<<<<< Updated upstream
const taskCols = `id, description, goal, exploration_id, status, paused, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
=======
func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
if len(companyIDs) == 0 {
return nil
}
var inserted int
err := tx.QueryRow(`
WITH requested(company_id, position) AS (
SELECT company_id, position
FROM unnest($2::bigint[]) WITH ORDINALITY AS requested(company_id, position)
), inserted AS (
INSERT INTO task_scope(task_id, kind, company_id, source, reason)
SELECT $1, 'company', companies.id, 'manual', '任务创建时关联企业'
FROM requested
JOIN companies ON companies.id=requested.company_id
ORDER BY requested.position
RETURNING company_id
)
SELECT count(*) FROM inserted`, taskID, companyIDs).Scan(&inserted)
if err != nil {
return err
}
if inserted != len(companyIDs) {
return fmt.Errorf("%w: one or more companies do not exist", ErrTaskCompanyNotFound)
}
return nil
}
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, COALESCE(name,''), description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Name, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil {
>>>>>>> Stashed changes
return nil, err
}
return &t, nil