同步
This commit is contained in:
1 parent
e67ff26e90
commit
58f23a2773
61 files changed
+11611
-443
No files matched your search
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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. 漏洞发现持久化
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
Reference in new issue
Block a user