feat(tasks): add templates findings workflows and batch controls

This commit is contained in:
pphua committed 2026-08-21 21:01:39 +08:00
1 parent b7e0b93023
commit d3183dbbd2
65 files changed
+11603 -2198

No files matched your search

+16 -1
View File
@@ -457,6 +457,18 @@ func isIdentChar(c byte) bool {
// ── QueryDSL ─────────────────────────────────────────────────────────────────
// ValidateDSL parses and compiles a DSL expression without touching the
// database. HTTP callers use it to distinguish client syntax errors from query
// failures, which must remain server errors.
func ValidateDSL(dsl string) error {
node, err := ParseDSL(dsl)
if err != nil {
return err
}
_, _, err = buildDSLWhere(node)
return err
}
// CountDSL returns the total number of assets matching a DSL expression (and optional
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
@@ -483,6 +495,9 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
node, err := ParseDSL(dsl)
if err != nil {
return nil, err
@@ -497,7 +512,7 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
}
args = append(args, limit, offset)
q := assetSelectCols + " WHERE " + where +
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
fmt.Sprintf(" ORDER BY last_seen DESC, id DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
+195 -45
View File
@@ -83,6 +83,7 @@ type PortService struct {
type AssetStore struct {
db *DB
company *CompanyStore
tx *sql.Tx
}
// Assets returns the asset store.
@@ -93,6 +94,39 @@ func (d *DB) Assets() *AssetStore {
// Companies returns the company store associated with this asset store.
func (s *AssetStore) Companies() *CompanyStore { return s.company }
// withCompanyScopeMutation serializes scope resolution and every asset write
// that consumes its result in one transaction. Nested asset side effects reuse
// the same transaction through the scoped store.
func (s *AssetStore) withCompanyScopeMutation(fn func(*AssetStore) (int64, error)) (int64, error) {
if s.tx != nil {
return fn(s)
}
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
id, err := fn(scoped)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return id, nil
}
func (s *AssetStore) resolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
if s.tx != nil {
return resolveCompanyWithICP(s.tx, rootDomain, ipStr, icp)
}
return s.company.ResolveCompanyWithICP(rootDomain, ipStr, icp)
}
// =====================================================================
// Helpers
// =====================================================================
@@ -219,7 +253,12 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
companyID, err := s.company.ResolveCompany(domain, "")
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertRootDomain(req)
})
}
companyID, err := s.resolveCompanyWithICP(domain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -237,12 +276,20 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, $4::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, company_source, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, 'scope', $4::bigint[])
ON CONFLICT (domain) WHERE type = 'root_domain' DO UPDATE SET
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
extra = assets.extra || EXCLUDED.extra,
last_seen = now()
@@ -267,8 +314,13 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
if req.IP == "" {
return 0, fmt.Errorf("ip is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertIP(req)
})
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany("", req.IP)
companyID, err := s.resolveCompanyWithICP("", req.IP, "")
if err != nil {
return 0, err
}
@@ -292,9 +344,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, company_source, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, 'scope', $6::bigint[])
ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
bound_domains = (SELECT ARRAY(SELECT DISTINCT unnest(assets.bound_domains || EXCLUDED.bound_domains))),
open_ports = (
@@ -306,7 +358,15 @@ ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
)
),
c_segment = COALESCE(assets.c_segment, EXCLUDED.c_segment),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.IP, csegVal, boundDomains, openPortsJSON, companyID, taskIDs).Scan(&id)
@@ -367,14 +427,19 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertSubdomain(req)
})
}
rootDomain, _ := RootDomain(domain)
if rootDomain == "" {
rootDomain = domain
}
// resolve company by root domain
companyID, err := s.company.ResolveCompany(rootDomain, "")
// Resolve by root domain first, then the asset's exact normalized ICP.
companyID, err := s.resolveCompanyWithICP(rootDomain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -409,7 +474,7 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
cseg := calcCSegment(ipStr)
if companyID == nil && ipStr != "" {
companyID, _ = s.company.ResolveCompany("", ipStr)
companyID, _ = s.resolveCompanyWithICP("", ipStr, "")
}
var taskIDs string
@@ -432,14 +497,22 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
recordType := req.RecordType
recordValueArr := marshalStringArray(req.RecordValue)
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, $9::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, company_source, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, 'scope', $9::bigint[])
ON CONFLICT (domain, COALESCE(record_type,'')) WHERE type = 'subdomain' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
record_value = (SELECT ARRAY(SELECT DISTINCT unnest(assets.record_value || EXCLUDED.record_value))),
extra = assets.extra || EXCLUDED.extra,
@@ -459,7 +532,7 @@ type UpsertAppReq struct {
Category string
Description string
ICP string
CompanyID *int64 // explicit override; nil = no auto attribution for apps
CompanyID *int64 // explicit override; nil = exact ICP auto-attribution when available
TaskID int64
}
@@ -468,6 +541,11 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
if req.Name == "" {
return 0, fmt.Errorf("app name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertApp(req)
})
}
var taskIDs string
if req.TaskID > 0 {
@@ -477,6 +555,7 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
var bundleVal, catVal, descVal, icpVal, companyIDVal any
companySource := "scope"
if req.BundleID != "" {
bundleVal = req.BundleID
}
@@ -491,36 +570,65 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
if req.CompanyID != nil {
companyIDVal = *req.CompanyID
companySource = "explicit"
} else {
companyID, err := s.resolveCompanyWithICP("", "", req.ICP)
if err != nil {
return 0, err
}
if companyID != nil {
companyIDVal = *companyID
}
}
var id int64
var err error
if req.BundleID != "" {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7, $8::bigint[])
ON CONFLICT (bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL DO UPDATE SET
app_name = COALESCE(EXCLUDED.app_name, assets.app_name),
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
} else {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6, $7::bigint[])
ON CONFLICT (app_name) WHERE type = 'app' AND bundle_id IS NULL DO UPDATE SET
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
}
return id, err
}
@@ -548,6 +656,11 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
if req.URL == "" {
return 0, fmt.Errorf("url is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertHTTPService(req)
})
}
normURL := normalizeURL(req.URL)
domain, port, serviceName := parseURL(normURL)
rootDomain, _ := RootDomain(domain)
@@ -556,7 +669,7 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -592,16 +705,16 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, service_type, service_name, domain, ip, port, root_domain,
c_segment, favicon_mmh3, technologies, status_code, content_length, page_title,
auth, company_id, task_ids
auth, company_id, company_source, task_ids
)
VALUES (
'service', $1, 'http', $2, $3, $4, $5, $6,
$7::cidr, $8, $9::text[], $10, $11, $12,
$13::jsonb[], $14, $15::bigint[]
$13::jsonb[], $14, 'scope', $15::bigint[]
)
ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
status_code = COALESCE(EXCLUDED.status_code, assets.status_code),
@@ -618,7 +731,15 @@ ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, serviceName, domainVal, ipVal, nullableInt(port), rootDomain,
@@ -675,6 +796,11 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
if req.ServiceName == "" {
return 0, fmt.Errorf("service_name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertOtherService(req)
})
}
// normalize service_name (lowercase) so the (domain,ip,port,service_name)
// dedup key doesn't split "SSH" and "ssh" into separate rows.
serviceName := strings.ToLower(strings.TrimSpace(req.ServiceName))
@@ -689,7 +815,7 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -721,12 +847,12 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, service_type, service_name, domain, ip, port, root_domain,
c_segment, auth, company_id, task_ids
c_segment, auth, company_id, company_source, task_ids
)
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, $9::bigint[])
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, 'scope', $9::bigint[])
ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other' DO UPDATE SET
domain = COALESCE(EXCLUDED.domain, assets.domain),
ip = COALESCE(EXCLUDED.ip, assets.ip),
@@ -738,7 +864,15 @@ ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE typ
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
serviceName, domainVal, ipVal, req.Port, rootDomainVal,
@@ -799,6 +933,11 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
if req.Method == "" {
return 0, fmt.Errorf("method is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertEndpoint(req)
})
}
method := strings.ToUpper(req.Method)
normURL := normalizeURL(req.URL)
domain, port, _ := parseURL(normURL)
@@ -808,7 +947,7 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -840,12 +979,12 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, method, domain, ip, port, root_domain,
params, company_id, task_ids, c_segment
params, company_id, company_source, task_ids, c_segment
)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, $9::bigint[], $10::cidr)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, 'scope', $9::bigint[], $10::cidr)
ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
params = (
SELECT ARRAY(
@@ -857,7 +996,15 @@ ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, method, domainVal, ipVal, nullableInt(port), rootDomainVal,
@@ -884,6 +1031,9 @@ func (s *AssetStore) QueryByType(typ string, limit, offset int) ([]*Asset, error
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
rows, err := s.db.Query(`
SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
@@ -897,7 +1047,7 @@ SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(method,''), array_to_json(params)::text, extra, last_seen::text
FROM assets
WHERE type = $1
ORDER BY last_seen DESC
ORDER BY last_seen DESC, id DESC
LIMIT $2 OFFSET $3`, typ, limit, offset)
if err != nil {
return nil, err
@@ -954,7 +1104,7 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) {
}
func pageClause(args *[]any, limit, offset int) string {
q := ` ORDER BY last_seen DESC`
q := ` ORDER BY last_seen DESC, id DESC`
if limit > 0 {
*args = append(*args, limit)
q += fmt.Sprintf(` LIMIT $%d`, len(*args))
+50
View File
@@ -4,6 +4,7 @@ import (
"fmt"
"slices"
"testing"
"time"
)
// testSetup opens a DB and returns both stores. Skips if no PG.
@@ -896,3 +897,52 @@ func TestNormalizeURL(t *testing.T) {
}
}
}
func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
d, assets, _ := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
taskID := stamp
marker := fmt.Sprintf("stable-page-%d", stamp)
sharedSeen := time.Date(2026, time.August, 21, 9, 0, 0, 0, time.UTC)
ids := make([]int64, 0, 5)
for i := 0; i < 5; i++ {
var id int64
domain := fmt.Sprintf("%s-%d.invalid", marker, i)
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,task_ids,last_seen)
VALUES ('root_domain',$1,$1,ARRAY[$2]::bigint[],$3) RETURNING id`,
domain, taskID, sharedSeen).Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
defer func() { _, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[])`, ids) }()
want := slices.Clone(ids)
slices.Reverse(want)
collect := func(query func(limit, offset int) ([]*Asset, error)) []int64 {
t.Helper()
var got []int64
for offset := 0; offset < len(ids); offset += 2 {
page, err := query(2, offset)
if err != nil {
t.Fatal(err)
}
for _, asset := range page {
got = append(got, asset.ID)
}
}
return got
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryByTask(taskID, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("task pagination order=%v want=%v", got, want)
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryDSL(marker, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("DSL pagination order=%v want=%v", got, want)
}
}
+457 -109
View File
@@ -2,9 +2,11 @@ package db
import (
"database/sql"
"errors"
"fmt"
"net"
"strings"
"unicode/utf8"
)
// =====================================================================
@@ -25,7 +27,7 @@ type Company struct {
type CompanyWithScope struct {
Company
Scope []ScopeRule `json:"scope"`
AssetCount int `json:"asset_count"`
AssetCount int `json:"asset_count"`
}
// ScopeRule is one company_scope row.
@@ -35,6 +37,7 @@ type ScopeRule struct {
Kind string `json:"kind"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Value string `json:"value,omitempty"`
Raw string `json:"raw"`
Reason string `json:"reason,omitempty"`
}
@@ -42,6 +45,51 @@ type ScopeRule struct {
// CompanyStore operates on the companies + company_scope tables.
type CompanyStore struct{ db *DB }
var (
ErrCompanyNameConflict = errors.New("company name already exists")
ErrCompanyNotFound = errors.New("company not found")
)
const (
// MaxCompanyScopeRules bounds both request processing and the total number
// of persisted rules for one company.
MaxCompanyScopeRules = 256
// Raw and normalized textual scope payloads are bounded by Unicode rune
// count so multi-byte input is treated consistently by the API and DB layer.
MaxCompanyScopeRawRunes = 1024
MaxCompanyScopeValueRunes = 1024
)
// CompanyScopeValidationError identifies a client-correctable scope error.
// Storage and transaction failures are returned as ordinary errors instead.
type CompanyScopeValidationError struct{ Message string }
func (e *CompanyScopeValidationError) Error() string { return e.Message }
// ValidateCompanyScopeInputBounds applies request-wide limits before parsing.
// Store methods call it again so non-HTTP callers cannot bypass the limits.
func ValidateCompanyScopeInputBounds(inputs []ScopeInput) error {
if len(inputs) > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 最多 %d 条", MaxCompanyScopeRules,
)}
}
for i, input := range inputs {
if utf8.RuneCountInString(input.Value) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
}
return nil
}
// Scope writes rebuild derived asset ownership globally, so serialize them to
// ensure the committed attribution always reflects the latest committed rules.
// This key is reserved for company mutations; 7337741001 is the schema lock and
// 7337741002 is the cross-package test-suite lock.
const companyScopeMutationLock int64 = 7337741003
// Companies returns the company store.
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
@@ -69,6 +117,60 @@ RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
return
}
// CreateCompanyWithScope creates a company without updating an existing row.
// The company, its valid initial scope rules, and derived asset attribution are
// committed atomically. Invalid inputs retain the legacy partial-validation
// contract and are reported without preventing valid rules from being stored.
func (s *CompanyStore) CreateCompanyWithScope(name, logo string, inputs []ScopeInput, reason string) (
id int64, added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, 0, nil, err
}
rules, invalid, validationErrors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
nkey := companyNKey(name)
var logoVal any
if logo != "" {
logoVal = logo
}
if err := tx.QueryRow(`
INSERT INTO companies(name, nkey, logo)
VALUES ($1, $2, $3)
ON CONFLICT (nkey) DO NOTHING
RETURNING id`, name, nkey, logoVal).Scan(&id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return 0, 0, 0, invalid, validationErrors, ErrCompanyNameConflict
}
return 0, 0, 0, invalid, validationErrors, err
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, id, rules, reason)
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
}
if err := tx.Commit(); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
return id, added, skipped, invalid, validationErrors, nil
}
// GetCompany returns one company by id (nil if not found).
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
c := &Company{}
@@ -102,12 +204,55 @@ func (s *CompanyStore) UpsertByName(name string) (int64, error) {
return id, err
}
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
// DeleteCompany deletes a company and re-evaluates automatic ownership against
// the remaining companies in the same transaction. Explicitly-owned assets are
// detached by the FK and may then fall back to a remaining scope match.
func (s *CompanyStore) DeleteCompany(id int64) error {
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
_, err := s.DeleteCompanyWithAssets(id, false)
return err
}
// DeleteCompanyWithAssets deletes a company and optionally all of its assets in
// one transaction, then re-evaluates ownership against the remaining companies.
func (s *CompanyStore) DeleteCompanyWithAssets(id int64, deleteAssets bool) (assetsDeleted int64, err error) {
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
if deleteAssets {
res, err := tx.Exec(`DELETE FROM assets WHERE company_id = $1`, id)
if err != nil {
return 0, err
}
assetsDeleted, err = res.RowsAffected()
if err != nil {
return 0, err
}
}
res, err := tx.Exec(`DELETE FROM companies WHERE id = $1`, id)
if err != nil {
return 0, err
}
deleted, err := res.RowsAffected()
if err != nil {
return 0, err
}
if deleted == 0 {
return 0, ErrCompanyNotFound
}
if err := recomputeAttributionTx(tx); err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return assetsDeleted, nil
}
// ListCompanies returns all companies with scope and asset count.
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
rows, err := s.db.Query(`
@@ -149,7 +294,7 @@ ORDER BY c.name`)
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
rows, err := s.db.Query(`
SELECT id, company_id, kind,
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
COALESCE(domain,''), COALESCE(net::text,''), COALESCE(value,''), raw, COALESCE(reason,'')
FROM company_scope
WHERE company_id = $1
ORDER BY id`, companyID)
@@ -157,10 +302,10 @@ ORDER BY id`, companyID)
return nil, err
}
defer rows.Close()
var out []ScopeRule
out := make([]ScopeRule, 0)
for rows.Next() {
var r ScopeRule
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Value, &r.Raw, &r.Reason); err != nil {
return nil, err
}
out = append(out, r)
@@ -171,48 +316,172 @@ ORDER BY id`, companyID)
// AddScope parses and inserts scope lines for a company, then reattributes assets.
// Returns counts of added, skipped, and invalid lines.
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
rule, err := ParseScopeLine(line)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
inserted, err := s.insertScopeRule(companyID, rule, reason)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
if inserted {
added++
} else {
skipped++
}
inputs = append(inputs, ScopeInput{Value: line})
}
// attribute newly-covered assets
if added > 0 {
_ = s.attributeByCompany(companyID)
}
return
return s.AddScopeInputs(companyID, inputs, reason)
}
// insertScopeRule inserts a scope rule. Returns (inserted, error):
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
// AddScopeInputs inserts structured scope rules. Empty kinds retain the legacy
// domain/IP/CIDR auto-detection used by AddScope.
func (s *CompanyStore) AddScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, skipped, invalid int, errors []string) {
added, skipped, invalid, validationErrors, err := s.AddScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, skipped, invalid, validationErrors
}
// AddScopeInputsChecked inserts structured scope rules while keeping input
// validation separate from storage and transaction errors.
func (s *CompanyStore) AddScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, nil, err
}
rules, invalid, errors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, invalid, errors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, invalid, errors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if len(rules) == 0 {
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return 0, 0, invalid, errors, nil
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, invalid, errors, fmt.Errorf("重新计算企业归属失败: %w", err)
}
}
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return added, skipped, invalid, errors, nil
}
func parseScopeInputs(inputs []ScopeInput) (rules []ParsedScope, invalid int, validationErrors []string) {
rules = make([]ParsedScope, 0, len(inputs))
for _, input := range inputs {
rule, err := ParseScopeInput(input)
if err != nil {
invalid++
validationErrors = append(validationErrors, fmt.Sprintf("%s: %v", input.Value, err))
continue
}
rules = append(rules, rule)
}
return rules, invalid, validationErrors
}
func validateParsedScopeBounds(rules []ParsedScope) error {
for i, rule := range rules {
if utf8.RuneCountInString(rule.Raw) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
if utf8.RuneCountInString(rule.Value) > MaxCompanyScopeValueRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条规范化值过长: 最多 %d 个字符", i+1, MaxCompanyScopeValueRunes,
)}
}
}
return nil
}
func ensureCompanyExistsTx(tx *sql.Tx, companyID int64) error {
var exists bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM companies WHERE id = $1)`, companyID).Scan(&exists); err != nil {
return err
}
if !exists {
return ErrCompanyNotFound
}
return nil
}
func ensureCompanyScopeCapacityTx(tx *sql.Tx, companyID int64) error {
var count int
if err := tx.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id = $1`, companyID).Scan(&count); err != nil {
return err
}
if count > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 每个企业最多 %d 条", MaxCompanyScopeRules,
)}
}
return nil
}
func lockCompanyScopeMutation(tx *sql.Tx) error {
_, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, companyScopeMutationLock)
return err
}
func insertScopeRulesTx(tx *sql.Tx, companyID int64, rules []ParsedScope, reason string) (
added, skipped int, needsAttribution bool, err error,
) {
for _, rule := range rules {
inserted, insertErr := insertScopeRuleTx(tx, companyID, rule, reason)
if insertErr != nil {
return 0, 0, false, insertErr
}
if !inserted {
skipped++
continue
}
added++
needsAttribution = needsAttribution || rule.Kind != "keyword"
}
return added, skipped, needsAttribution, nil
}
// insertScopeRuleTx inserts a scope rule. inserted=false means a duplicate was
// ignored by ON CONFLICT, not an error.
func insertScopeRuleTx(tx *sql.Tx, companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
var res interface{ RowsAffected() (int64, error) }
if rule.Kind == "domain" {
res, err = s.db.Exec(`
switch rule.Kind {
case "domain":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, domain, raw, reason)
VALUES ($1, 'domain', $2, $3, $4)
ON CONFLICT ON CONSTRAINT uq_sv2_domain DO NOTHING`,
companyID, rule.Domain, rule.Raw, reason)
} else {
res, err = s.db.Exec(`
case "ip", "cidr":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, net, raw, reason)
VALUES ($1, $2, $3::cidr, $4, $5)
ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
companyID, rule.Kind, rule.Net, rule.Raw, reason)
case "icp", "keyword":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, value, raw, reason)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (company_id, kind, value) WHERE kind IN ('icp','keyword') DO NOTHING`,
companyID, rule.Kind, rule.Value, rule.Raw, reason)
default:
return false, fmt.Errorf("unsupported company scope kind %q", rule.Kind)
}
if err != nil {
return false, err
@@ -221,114 +490,179 @@ ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
return n > 0, nil
}
// RecomputeAttribution clears all company attribution on assets and rebuilds
// from scope rules. domain rules use root_domain equality; ip/cidr use INET >>= .
// RecomputeAttribution rebuilds only scope-derived ownership. Explicit company
// links are immutable under scope edits. Precedence is domain, IP/CIDR, then
// normalized exact ICP; keyword rules never attribute assets.
func (s *CompanyStore) RecomputeAttribution() error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// step 1: clear all attribution
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL`); err != nil {
if err := lockCompanyScopeMutation(tx); err != nil {
return err
}
// step 2: domain-based attribution (root_domain exact match)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind = 'domain'
AND a.root_domain = cs.domain
ORDER BY length(cs.domain) DESC
LIMIT 1
)
WHERE type IN ('root_domain','subdomain','service','endpoint')
AND root_domain IS NOT NULL`); err != nil {
if err := recomputeAttributionTx(tx); err != nil {
return err
}
// step 3: ip/cidr attribution (only for assets not yet attributed)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
ORDER BY masklen(cs.net) DESC
LIMIT 1
)
WHERE type IN ('ip','subdomain','service','endpoint')
AND ip IS NOT NULL
AND company_id IS NULL`); err != nil {
return err
}
return tx.Commit()
}
// attributeByCompany reattributes only assets matching a specific company's scope.
// Cheaper than full recompute; used when a company's scope is updated.
func (s *CompanyStore) attributeByCompany(companyID int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// clear this company's existing claims
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL WHERE company_id = $1`, companyID); err != nil {
return err
}
// domain rules
func recomputeAttributionTx(tx *sql.Tx) error {
// Only derived rows are cleared. Historical rows migrated without provenance
// are marked explicit by schema.sql, which is the non-destructive default.
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind = 'domain'
AND a.root_domain = cs.domain
AND type IN ('root_domain','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
UPDATE assets
SET company_id = NULL, company_source = 'scope'
WHERE company_source = 'scope'`); err != nil {
return err
}
// ip/cidr rules
// Domain-based attribution (root_domain exact match).
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
AND type IN ('ip','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'domain' AND a.root_domain = cs.domain
WHERE a.company_id IS NULL
AND a.type IN ('root_domain','subdomain','service','endpoint')
AND a.root_domain IS NOT NULL
ORDER BY a.id, length(cs.domain) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return tx.Commit()
// IP/CIDR attribution for still-unowned assets.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind IN ('ip','cidr') AND cs.net >>= a.ip::inet
WHERE a.company_id IS NULL
AND a.type IN ('ip','subdomain','service','endpoint')
AND a.ip IS NOT NULL
ORDER BY a.id, masklen(cs.net) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
// Exact normalized ICP attribution after domain/network precedence.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'icp'
AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = cs.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = cs.value
)
WHERE a.company_id IS NULL
AND (COALESCE(a.icp,'') <> '' OR COALESCE(a.app_icp,'') <> '')
ORDER BY a.id, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return nil
}
// UpdateScope replaces all scope rules for a company and reattributes.
func (s *CompanyStore) UpdateScope(companyID int64, lines []string, reason string) (added, invalid int, errs []string) {
// delete existing rules
if _, err := s.db.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
errs = append(errs, err.Error())
return
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
inputs = append(inputs, ScopeInput{Value: line})
}
added, _, invalid, errs = s.AddScope(companyID, lines, reason)
return
return s.UpdateScopeInputs(companyID, inputs, reason)
}
// UpdateScopeInputs replaces all rules with a structured set.
func (s *CompanyStore) UpdateScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, invalid int, errs []string) {
added, invalid, validationErrors, err := s.UpdateScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, invalid, validationErrors
}
// UpdateScopeInputsChecked replaces all rules while separating validation
// feedback from storage and transaction failures.
func (s *CompanyStore) UpdateScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, nil, err
}
rules, invalid, errs := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, invalid, errs, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, invalid, errs, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
if _, err := tx.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
return 0, invalid, errs, err
}
added, _, _, err = insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
// Rebuild even for an empty replacement because removing the old rules may
// detach scope-derived assets or expose a lower-precedence company match.
if err := recomputeAttributionTx(tx); err != nil {
return 0, invalid, errs, fmt.Errorf("重新计算企业归属失败: %w", err)
}
if err := tx.Commit(); err != nil {
return 0, invalid, errs, err
}
return added, invalid, errs, nil
}
// ResolveCompany returns the company_id for a given root_domain and/or ip, or nil
// if no scope rule matches. Mirrors the attribution logic used at asset insert time.
func (s *CompanyStore) ResolveCompany(rootDomain, ipStr string) (*int64, error) {
return s.ResolveCompanyWithICP(rootDomain, ipStr, "")
}
// ResolveCompanyWithICP mirrors RecomputeAttribution for insert-time ownership.
// ICP is consulted only after domain and IP/CIDR fail to match.
func (s *CompanyStore) ResolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
return resolveCompanyWithICP(s.db, rootDomain, ipStr, icp)
}
type companyScopeQueryer interface {
QueryRow(query string, args ...any) *sql.Row
}
func resolveCompanyWithICP(q companyScopeQueryer, rootDomain, ipStr, icp string) (*int64, error) {
if rootDomain != "" {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'domain'
AND domain = $1
ORDER BY length(domain) DESC
ORDER BY length(domain) DESC, company_id
LIMIT 1`, rootDomain).Scan(&cid)
if err == nil {
return &cid, nil
@@ -340,11 +674,11 @@ LIMIT 1`, rootDomain).Scan(&cid)
if ipStr != "" {
if net.ParseIP(ipStr) != nil {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind IN ('ip','cidr')
AND net >>= $1::inet
ORDER BY masklen(net) DESC
ORDER BY masklen(net) DESC, company_id
LIMIT 1`, ipStr).Scan(&cid)
if err == nil {
return &cid, nil
@@ -354,5 +688,19 @@ LIMIT 1`, ipStr).Scan(&cid)
}
}
}
if normalized := NormalizeICP(icp); normalized != "" {
var cid int64
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'icp' AND value = $1
ORDER BY company_id
LIMIT 1`, normalized).Scan(&cid)
if err == nil {
return &cid, nil
}
if err != sql.ErrNoRows {
return nil, err
}
}
return nil, nil
}
+150
View File
@@ -1,7 +1,11 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
// cleanup helpers to remove test data
@@ -57,6 +61,70 @@ func TestCompanyUpsertAndGet(t *testing.T) {
}
}
func TestCreateCompanyWithScopeRejectsNormalizedDuplicate(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("Strict Company %d", stamp)
id, added, _, _, validationErrors, err := cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("strict-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
if added != 1 || len(validationErrors) != 0 {
t.Fatalf("initial scope: added=%d errors=%v", added, validationErrors)
}
_, _, _, _, _, err = cs.CreateCompanyWithScope(
" "+strings.ToUpper(strings.ReplaceAll(name, " ", " "))+" ",
"",
[]ScopeInput{{Kind: "domain", Value: fmt.Sprintf("replacement-%d.example", stamp)}},
"test",
)
if !errors.Is(err, ErrCompanyNameConflict) {
t.Fatalf("duplicate create error=%v want ErrCompanyNameConflict", err)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("strict-%d.example", stamp) {
t.Fatalf("duplicate create changed existing scope: %+v", scope)
}
}
func TestCreateCompanyWithScopeRollsBackOnScopeWriteFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Create %d", time.Now().UnixNano())
_, _, _, _, _, err = cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "test")
if err == nil {
t.Fatal("expected scope database write to fail")
}
company, getErr := cs.GetCompanyByName(name)
if getErr != nil {
t.Fatal(getErr)
}
if company != nil {
defer cleanupCompany(d, company.ID)
t.Fatalf("company row survived failed initial scope transaction: %+v", company)
}
}
func TestCompanyScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -189,6 +257,41 @@ func TestUpdateScope(t *testing.T) {
}
}
func TestUpdateScopeRollsBackOnInsertFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Scope Update %d", time.Now().UnixNano())
id, _, err := cs.UpsertCompany(name, "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
oldDomain := fmt.Sprintf("old-%d.example", time.Now().UnixNano())
added, _, invalid, addErrors := cs.AddScopeInputs(id, []ScopeInput{{Kind: "domain", Value: oldDomain}}, "initial")
if added != 1 || invalid != 0 || len(addErrors) != 0 {
t.Fatalf("seed scope: added=%d invalid=%d errors=%v", added, invalid, addErrors)
}
added, invalid, updateErrors := cs.UpdateScopeInputs(id, []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "replacement")
if added != 0 || invalid != 0 || len(updateErrors) == 0 {
t.Fatalf("failed update result: added=%d invalid=%d errors=%v", added, invalid, updateErrors)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != oldDomain {
t.Fatalf("failed replacement did not preserve old scope: %+v", scope)
}
}
func TestDeleteCompany(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -217,6 +320,53 @@ func TestDeleteCompany(t *testing.T) {
}
}
func TestDeleteCompanyWithAssetsDeletesBoth(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
id, _, err := cs.UpsertCompany(fmt.Sprintf("Delete Assets Company %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
var assetID int64
domain := fmt.Sprintf("delete-assets-%d.example", stamp)
if err := d.QueryRow(`
INSERT INTO assets(type, domain, root_domain, company_id, company_source)
VALUES ('root_domain', $1, $1, $2, 'explicit')
RETURNING id`, domain, id).Scan(&assetID); err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id = $1`, assetID) //nolint:errcheck
deleted, err := cs.DeleteCompanyWithAssets(id, true)
if err != nil {
t.Fatal(err)
}
if deleted != 1 {
t.Fatalf("assets deleted=%d want 1", deleted)
}
company, err := cs.GetCompany(id)
if err != nil {
t.Fatal(err)
}
if company != nil {
t.Fatalf("company still exists: %+v", company)
}
var assetsRemaining int
if err := d.QueryRow(`SELECT COUNT(*) FROM assets WHERE id = $1`, assetID).Scan(&assetsRemaining); err != nil {
t.Fatal(err)
}
if assetsRemaining != 0 {
t.Fatalf("asset %d survived company deletion", assetID)
}
}
func TestRecomputeAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
+15
View File
@@ -0,0 +1,15 @@
package db
import "testing"
func TestCompanyScopeMutationLockDoesNotReuseInfrastructureLocks(t *testing.T) {
reserved := map[string]int64{
"schema migration": 7337741001,
"cross-package test suite": 7337741002,
}
for name, key := range reserved {
if companyScopeMutationLock == key {
t.Fatalf("company scope mutation lock reuses the %s advisory lock key %d", name, key)
}
}
}
+61 -2
View File
@@ -4,19 +4,78 @@ import (
"fmt"
"net"
"strings"
"unicode"
"golang.org/x/net/publicsuffix"
)
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
// ParsedScope is one parsed asset-scope entry. Its kind selects Domain, Net, or Value.
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
type ParsedScope struct {
Kind string // "domain" | "ip" | "cidr"
Kind string // "domain" | "ip" | "cidr" | "icp" | "keyword"
Domain string // normalized registrable/root domain (kind=domain)
Net string // normalized CIDR, single IP as /32 or /128 (kind=ip|cidr)
Value string // normalized text (kind=icp|keyword)
Raw string // original input line
}
// ScopeInput is the structured API form for a company scope rule. Empty Kind
// preserves the legacy one-line auto detection for domain/IP/CIDR clients.
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), " "))
}
// ParseScopeInput validates an explicitly typed rule. Legacy callers can omit
// Kind and retain the existing domain/IP/CIDR auto-detection behavior.
func ParseScopeInput(input ScopeInput) (ParsedScope, error) {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
raw := strings.TrimSpace(input.Value)
if kind == "" {
return ParseScopeLine(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)
}
}
// 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.
+248
View File
@@ -0,0 +1,248 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
domain := fmt.Sprintf("scope-lock-%d.example", stamp)
oldCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock Old %d", stamp), "")
if err != nil {
t.Fatal(err)
}
newCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock New %d", stamp), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain = $1`, domain)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
})
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(oldCompany, []ScopeInput{
{Kind: "domain", Value: domain},
}, "test"); err != nil || added != 1 || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("seed old scope: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
mutation, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer mutation.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(mutation); err != nil {
t.Fatal(err)
}
if _, err := mutation.Exec(`DELETE FROM company_scope WHERE company_id=$1 AND kind='domain' AND domain=$2`, oldCompany, domain); err != nil {
t.Fatal(err)
}
if _, err := insertScopeRuleTx(mutation, newCompany, ParsedScope{
Kind: "domain", Domain: domain, Raw: domain,
}, "test"); err != nil {
t.Fatal(err)
}
started := make(chan struct{})
type result struct {
id int64
err error
}
resultCh := make(chan result, 1)
go func() {
close(started)
id, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
resultCh <- result{id: id, err: err}
}()
<-started
select {
case got := <-resultCh:
t.Fatalf("asset upsert escaped scope mutation lock: %+v", got)
case <-time.After(100 * time.Millisecond):
}
if err := mutation.Commit(); err != nil {
t.Fatal(err)
}
var got result
select {
case got = <-resultCh:
case <-time.After(3 * time.Second):
t.Fatal("asset upsert did not resume after scope mutation committed")
}
if got.err != nil {
t.Fatal(got.err)
}
var companyID *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, got.id).Scan(&companyID); err != nil {
t.Fatal(err)
}
if companyID == nil || *companyID != newCompany {
t.Fatalf("asset retained stale company: got=%v want=%d", companyID, newCompany)
}
}
func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
first, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break First %d", stamp), "")
if err != nil {
t.Fatal(err)
}
second, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break Second %d", stamp), "")
if err != nil {
t.Fatal(err)
}
want := min(first, second)
domain := fmt.Sprintf("tie-break-%d.example", stamp)
segment := uint64(stamp) & 0xffff
network := fmt.Sprintf("2001:db8:%x:1::/64", segment)
ip := fmt.Sprintf("2001:db8:%x:1::42", segment)
icp := fmt.Sprintf("ICP-TIE-%d", stamp)
appName := fmt.Sprintf("Tie Break App %d", stamp)
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
})
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "cidr", Value: network},
{Kind: "icp", Value: icp},
}
for _, companyID := range []int64{second, first} {
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != len(rules) || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("add company %d rules: added=%d invalid=%d validation=%v err=%v", companyID, added, invalid, validationErrors, err)
}
}
assertResolved := func(label string, got *int64, err error) {
t.Helper()
if err != nil {
t.Fatalf("%s resolve: %v", label, err)
}
if got == nil || *got != want {
t.Fatalf("%s resolve=%v want lowest company id %d", label, got, want)
}
}
got, err := companies.ResolveCompany(domain, "")
assertResolved("domain", got, err)
got, err = companies.ResolveCompany("", ip)
assertResolved("network", got, err)
got, err = companies.ResolveCompanyWithICP("", "", icp)
assertResolved("icp", got, err)
rootID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
if err != nil {
t.Fatal(err)
}
ipID, err := assets.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
appID, err := assets.UpsertApp(UpsertAppReq{Name: appName, ICP: icp})
if err != nil {
t.Fatal(err)
}
assertAssets := func(stage string) {
t.Helper()
for _, id := range []int64{rootID, ipID, appID} {
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, id).Scan(&got); err != nil {
t.Fatal(err)
}
if got == nil || *got != want {
t.Fatalf("%s asset %d company=%v want=%d", stage, id, got, want)
}
}
}
assertAssets("live")
if err := companies.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertAssets("recomputed")
}
func TestCompanyScopeLimitsAndCheckedErrors(t *testing.T) {
boundary := strings.Repeat("界", MaxCompanyScopeRawRunes)
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary}}); err != nil {
t.Fatalf("exact raw rune boundary rejected: %v", err)
}
var validationErr *CompanyScopeValidationError
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary + "界"}}); !errors.As(err, &validationErr) {
t.Fatalf("oversized raw value error=%v want CompanyScopeValidationError", err)
}
tooMany := make([]ScopeInput, MaxCompanyScopeRules+1)
if err := ValidateCompanyScopeInputBounds(tooMany); !errors.As(err, &validationErr) {
t.Fatalf("too many rules error=%v want CompanyScopeValidationError", err)
}
d, _, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Limits %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
rules := make([]ScopeInput, MaxCompanyScopeRules)
for i := range rules {
rules[i] = ScopeInput{Kind: "keyword", Value: fmt.Sprintf("limit-%d-%d", stamp, i)}
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != MaxCompanyScopeRules {
t.Fatalf("fill scope capacity: added=%d err=%v", added, err)
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, []ScopeInput{
{Kind: "keyword", Value: fmt.Sprintf("limit-%d-overflow", stamp)},
}, "test"); added != 0 || !errors.As(err, &validationErr) {
t.Fatalf("scope overflow: added=%d err=%v", added, err)
}
var count int
if err := d.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id=$1`, companyID).Scan(&count); err != nil {
t.Fatal(err)
}
if count != MaxCompanyScopeRules {
t.Fatalf("failed capacity mutation changed scope count=%d want=%d", count, MaxCompanyScopeRules)
}
missingID := companyID + 1_000_000_000
if _, _, _, _, err := companies.AddScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing add error=%v want ErrCompanyNotFound", err)
}
if _, _, _, err := companies.UpdateScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing update error=%v want ErrCompanyNotFound", err)
}
if _, err := companies.DeleteCompanyWithAssets(missingID, true); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing delete error=%v want ErrCompanyNotFound", err)
}
}
func TestCompanyScopeCheckedReturnsSystemErrorSeparately(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
companies := d.Companies()
if err := d.Close(); err != nil {
t.Fatal(err)
}
_, _, _, validationErrors, err := companies.AddScopeInputsChecked(1, nil, "test")
if err == nil {
t.Fatal("closed database did not return a system error")
}
if len(validationErrors) != 0 {
t.Fatalf("system error leaked into validation errors: %v", validationErrors)
}
var validationErr *CompanyScopeValidationError
if errors.As(err, &validationErr) || errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("system error misclassified: %v", err)
}
}
+195 -2
View File
@@ -3,6 +3,7 @@ package db
import (
"fmt"
"testing"
"time"
)
// TestParseScopeLine covers classification + guardrails without a DB.
@@ -15,9 +16,9 @@ func TestParseScopeLine(t *testing.T) {
{"example.com", "domain", false},
{"https://sub.example.com/path", "domain", false},
{"1.2.3.4", "ip", false},
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"198.51.100.0/24", "cidr", false},
{"co.uk", "", true}, // bare public suffix
{"co.uk", "", true}, // bare public suffix
{"not a host", "", true},
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
}
@@ -39,6 +40,198 @@ func TestParseScopeLine(t *testing.T) {
}
}
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)
}
// 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
+47 -11
View File
@@ -50,19 +50,32 @@ GROUP BY c.id`)
// Conversation is one ChatGPT-style chat thread bound to an agent key. It lives
// independent of the pentest exploration graph — see schema.sql §I.
type Conversation struct {
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
Pinned bool `json:"pinned"`
PinnedAt *time.Time `json:"pinned_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at`
// ConversationPatch updates only the fields whose pointers are non-nil.
type ConversationPatch struct {
Title *string
Pinned *bool
}
const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at`
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
var c Conversation
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt)
var pinnedAt sql.NullTime
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt)
if pinnedAt.Valid {
c.Pinned = true
c.PinnedAt = &pinnedAt.Time
}
return c, err
}
@@ -127,7 +140,8 @@ func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error {
// ListConversations returns all threads, most-recently-updated first.
func (d *DB) ListConversations() ([]*Conversation, error) {
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`)
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations
ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
@@ -155,9 +169,31 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) {
return &c, nil
}
// RenameConversation sets a thread's title.
// UpdateConversation applies a partial title/pin mutation and returns the updated
// row. Pinning an already-pinned conversation preserves its original pin order.
func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`UPDATE conversations SET
title = CASE WHEN $2::boolean THEN $3 ELSE title END,
pinned_at = CASE
WHEN $4::boolean IS NULL THEN pinned_at
WHEN $4::boolean THEN COALESCE(pinned_at, now())
ELSE NULL
END
WHERE id=$1
RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &c, nil
}
// RenameConversation sets a thread's title. Kept for automatic first-message
// titles and compatibility with existing callers.
func (d *DB) RenameConversation(id int64, title string) error {
_, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title)
_, err := d.UpdateConversation(id, ConversationPatch{Title: &title})
return err
}
+69
View File
@@ -0,0 +1,69 @@
package db
import (
"fmt"
"testing"
"time"
)
func TestConversationPinOrderingAndPatch(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
first, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-first-%d", suffix), nil)
if err != nil {
t.Fatal(err)
}
second, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-second-%d", suffix), nil)
if err != nil {
_ = d.DeleteConversation(first.ID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(first.ID)
_ = d.DeleteConversation(second.ID)
})
pinned := true
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || !first.Pinned || first.PinnedAt == nil {
t.Fatalf("pin first = %+v, %v", first, err)
}
firstPinnedAt := *first.PinnedAt
time.Sleep(5 * time.Millisecond)
second, err = d.UpdateConversation(second.ID, ConversationPatch{Pinned: &pinned})
if err != nil || second == nil || !second.Pinned {
t.Fatalf("pin second = %+v, %v", second, err)
}
newTitle := "renamed while pinned"
first, err = d.UpdateConversation(first.ID, ConversationPatch{Title: &newTitle, Pinned: &pinned})
if err != nil || first == nil {
t.Fatal(err)
}
if first.Title != newTitle || first.PinnedAt == nil || !first.PinnedAt.Equal(firstPinnedAt) {
t.Fatalf("repeat pin should preserve pin time: %+v (want %v)", first, firstPinnedAt)
}
conversations, err := d.ListConversations()
if err != nil {
t.Fatal(err)
}
positions := map[int64]int{}
for index, conversation := range conversations {
positions[conversation.ID] = index
}
if positions[second.ID] >= positions[first.ID] {
t.Fatalf("newer pin must sort first: second=%d first=%d", positions[second.ID], positions[first.ID])
}
pinned = false
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || first.Pinned || first.PinnedAt != nil {
t.Fatalf("unpin first = %+v, %v", first, err)
}
}
+215 -6
View File
@@ -3,11 +3,17 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"time"
)
// ErrIntentStateConflict marks a failed compare-and-set transition. Callers use
// it to distinguish an expected control race (HTTP 409) from a storage failure.
var ErrIntentStateConflict = errors.New("intent state changed concurrently")
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
@@ -70,6 +76,17 @@ type TokenUsage struct {
CacheWriteTokens int `json:"cache_write_tokens"`
}
// SessionTokenUsage is the authoritative completed-run total for one task UI
// session. Worker sessions are keyed by intent rather than the reusable work#N
// executor name, so retries and reassignment remain attached to the same row.
type SessionTokenUsage struct {
Session string `json:"session"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// DailyTokenBucket is one day's global token aggregate across all tasks.
type DailyTokenBucket struct {
Day string `json:"day"` // "YYYY-MM-DD"
@@ -264,6 +281,21 @@ WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, termin
return err
}
// 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.
@@ -295,7 +327,7 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, err
}
if state != "running" && state != "paused" {
return out, fmt.Errorf("intent state %s cannot be cancelled", state)
return out, fmt.Errorf("%w: intent state %s cannot be cancelled", ErrIntentStateConflict, state)
}
if err := tx.QueryRow(`SELECT
@@ -304,14 +336,33 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
FROM exploration_edges e
JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil {
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`, s.expID, id, RelYields); err != nil {
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)
@@ -319,15 +370,36 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
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')
)
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
@@ -339,6 +411,108 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
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`
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
@@ -761,6 +935,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
return out, rows.Err()
}
// TokenStatsBySession aggregates every persisted completed run, independent of
// activity-history pagination. Main/planner use fixed keys; Worker executions use
// their intent node id, even when a different work#N process handled a retry.
func (s *ExplorationStore) TokenStatsBySession() ([]SessionTokenUsage, error) {
rows, err := s.db.Query(`SELECT
CASE
WHEN COALESCE(a.worker,'')='mainagent' THEN 'main'
WHEN COALESCE(a.worker,'')='planner' THEN 'plan'
WHEN n.id IS NOT NULL THEN 'intent:' || n.id::text
ELSE ''
END AS session_key,
COALESCE(SUM(a.input_tokens),0), COALESCE(SUM(a.output_tokens),0),
COALESCE(SUM(a.cache_read_tokens),0), COALESCE(SUM(a.cache_write_tokens),0)
FROM activity a
LEFT JOIN exploration_nodes n
ON n.id=a.node_id AND n.exploration_id=a.exploration_id AND n.kind='intent'
WHERE a.exploration_id=$1 AND a.kind='result'
AND (COALESCE(a.worker,'') IN ('mainagent','planner') OR n.id IS NOT NULL)
GROUP BY session_key
ORDER BY session_key`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []SessionTokenUsage{}
for rows.Next() {
var u SessionTokenUsage
if err := rows.Scan(&u.Session, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out = append(out, u)
}
return out, rows.Err()
}
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
// Returns items and the new cursor (max id seen).
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
+130
View File
@@ -0,0 +1,130 @@
package db
import (
"fmt"
"testing"
)
func TestTokenStatsBySessionUsesCompleteIntentHistory(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("token sessions", "aggregate complete session history")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "worker session"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
factID, err := store.AddNode(KindFact, map[string]any{"summary": "not a worker"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
otherExpID, err := d.CreateExploration("other token sessions", "foreign intent must not match")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, otherExpID)
foreignIntentID, err := d.Exploration(otherExpID).AddIntent(map[string]any{"summary": "foreign worker"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
appendResult := func(nodeID *int64, worker string, input, output, read, write int) {
t.Helper()
if _, appendErr := store.AppendActivity(Activity{
NodeID: nodeID,
Worker: worker,
Kind: "result",
InputTokens: &input,
OutputTokens: &output,
CacheReadTokens: &read,
CacheWriteTokens: &write,
}); appendErr != nil {
t.Fatal(appendErr)
}
}
appendResult(nil, "mainagent", 11, 12, 13, 14)
appendResult(nil, "planner", 21, 22, 23, 24)
// More than ActivityPage's default limit proves this aggregate reads the
// persisted history directly instead of summing the currently loaded page.
const completedRuns = 205
for i := 0; i < completedRuns; i++ {
appendResult(&intentID, fmt.Sprintf("work#%d", i%3+1), 1, 2, 3, 4)
}
// All three rows remain part of the legacy worker/whole-task totals, but none
// represents a local Worker session and therefore none may create intent:*.
appendResult(&factID, "work#fact", 31, 32, 33, 34)
appendResult(nil, "work#missing-node", 41, 42, 43, 44)
appendResult(&foreignIntentID, "work#foreign", 51, 52, 53, 54)
ignoredInput := 1000
if _, err := store.AppendActivity(Activity{
NodeID: &intentID,
Worker: "work#1",
Kind: "usage",
InputTokens: &ignoredInput,
}); err != nil {
t.Fatal(err)
}
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
got := make(map[string]SessionTokenUsage, len(sessions))
for _, usage := range sessions {
got[usage.Session] = usage
}
if len(got) != 3 {
t.Fatalf("sessions = %+v, want only main, plan, and the local intent", sessions)
}
assertSessionTokens(t, got["main"], "main", 11, 12, 13, 14)
assertSessionTokens(t, got["plan"], "plan", 21, 22, 23, 24)
assertSessionTokens(t, got[fmt.Sprintf("intent:%d", intentID)], fmt.Sprintf("intent:%d", intentID),
completedRuns, completedRuns*2, completedRuns*3, completedRuns*4)
workers, err := store.TokenStatsByWorker()
if err != nil {
t.Fatal(err)
}
if len(workers) != 8 { // main, planner, 3 executors, and the 3 deliberately invalid session rows
t.Fatalf("legacy workers changed: got %d entries: %+v", len(workers), workers)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total,
11+21+completedRuns+31+41+51,
12+22+completedRuns*2+32+42+52,
13+23+completedRuns*3+33+43+53,
14+24+completedRuns*4+34+44+54)
}
func assertSessionTokens(t *testing.T, got SessionTokenUsage, session string, input, output, read, write int) {
t.Helper()
if got.Session != session || got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("session %q = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
session, got, input, output, read, write)
}
}
func assertTokenUsage(t *testing.T, got TokenUsage, input, output, read, write int) {
t.Helper()
if got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("token total = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
got, input, output, read, write)
}
}
+200 -4
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
@@ -147,6 +148,11 @@ type FindingFilter struct {
Sort string // "severity" | "time"
}
// FindingUnassignedTask is the task filter sentinel for findings whose task is
// absent. That includes rows created without a task and rows retained after their
// originating task was deleted (the findings FK is ON DELETE SET NULL).
const FindingUnassignedTask = "__unassigned__"
// where builds the WHERE clause (shared by the page and count queries) plus its
// positional args. Only equality filters, all parameterized.
func (f FindingFilter) where() (string, []any) {
@@ -163,7 +169,9 @@ func (f FindingFilter) where() (string, []any) {
add("status", f.Status)
add("vulnclass", f.VulnClass)
// task_id 是 bigint 列,按整数比较(不能走上面的文本 add);空/非法值忽略。
if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
if f.TaskID == FindingUnassignedTask {
conds = append(conds, "(f.task_id IS NULL OR t.id IS NULL)")
} else if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
args = append(args, tid)
conds = append(conds, fmt.Sprintf("f.task_id = $%d", len(args)))
}
@@ -185,14 +193,19 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
where, args := f.where()
var total int
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f`+where, args...).Scan(&total); err != nil {
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f LEFT JOIN tasks t ON f.task_id=t.id`+where, args...).Scan(&total); err != nil {
return nil, 0, err
}
// Avoid overflowing (page-1)*pageSize for an arbitrarily large page number.
// Once the requested page is beyond the exact count, no data query is needed.
if total == 0 || page > (total-1)/pageSize+1 {
return []*DBFinding{}, total, nil
}
order := "f.created_at DESC"
order := "f.created_at DESC, f.id DESC"
if f.Sort == "severity" {
// critical > high > medium > low > 其它, then newest first.
order = `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`
order = `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, f.id DESC`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
q := fmt.Sprintf(`
@@ -210,6 +223,189 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
return out, total, err
}
// 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"`
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.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.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.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 导出(导出当前筛选/全部)。结果按严重等级降序、
+298
View File
@@ -1,8 +1,12 @@
package db
import (
"encoding/json"
"errors"
"fmt"
"slices"
"testing"
"time"
)
// TestDeleteFinding verifies删除漏洞 removes both the findings row and its
@@ -183,3 +187,297 @@ func TestFindingsPageAndStats(t *testing.T) {
t.Fatalf("SetFindingSeverity(-1): want 0,nil got %d,%v", n, err)
}
}
func TestListFindingsPageUsesStableIDTieBreaker(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_stable_page_%d__", time.Now().UnixNano())
defer d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc)
ids := make([]int64, 0, 7)
for i := 0; i < 7; i++ {
id, addErr := d.AddFinding(0, 0, vc, "", SeverityHigh, fmt.Sprintf("finding %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
ids = append(ids, id)
}
sharedCreatedAt := time.Date(2026, time.August, 21, 8, 30, 0, 0, time.UTC)
if _, err := d.Exec(`UPDATE findings SET created_at=$1 WHERE vulnclass=$2`, sharedCreatedAt, vc); err != nil {
t.Fatal(err)
}
want := make([]int64, len(ids))
for i := range ids {
want[i] = ids[len(ids)-1-i]
}
for _, sort := range []string{"time", "severity"} {
t.Run(sort, func(t *testing.T) {
var got []int64
for page := 1; page <= 3; page++ {
items, total, pageErr := d.ListFindingsPage(FindingFilter{VulnClass: vc, Sort: sort}, page, 3)
if pageErr != nil {
t.Fatal(pageErr)
}
if total != len(ids) {
t.Fatalf("page %d total=%d, want %d", page, total, len(ids))
}
for _, item := range items {
got = append(got, item.ID)
}
}
if !slices.Equal(got, want) {
t.Fatalf("same-timestamp pagination was unstable: got=%v want=%v", got, want)
}
})
}
}
func TestFindingGroupsAndUnassignedPaging(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_groups_%d__", time.Now().UnixNano())
cleanupFindings := func() { _, _ = d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc) }
defer cleanupFindings()
taskA, err := d.CreateTask("group A", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskA.ID)
taskB, err := d.CreateTask("group B", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskB.ID)
seed := []struct {
taskID int64
severity string
status string
}{
{taskA.ID, SeverityCritical, FindingPending},
{taskA.ID, SeverityHigh, FindingResolved},
{taskB.ID, SeverityLow, FindingPending},
{0, SeverityMedium, FindingPending},
}
for i, item := range seed {
id, addErr := d.AddFinding(item.taskID, 0, vc, "", item.severity, fmt.Sprintf("summary %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
if _, setErr := d.SetFindingStatus(id, item.status); setErr != nil {
t.Fatalf("SetFindingStatus[%d]: %v", i, setErr)
}
}
// Retained findings from deleted tasks join the same bucket as findings that
// were created without any task.
if err := d.DeleteTask(taskB.ID); err != nil {
t.Fatal(err)
}
page, groupTotal, findingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc, Sort: "severity"}, 1, 1)
if err != nil {
t.Fatal(err)
}
if len(page) != 1 || groupTotal != 2 || findingTotal != 4 {
t.Fatalf("page totals: items=%d groups=%d findings=%d", len(page), groupTotal, findingTotal)
}
groups, _, _, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
var live, unassigned *FindingGroup
for i := range groups {
if groups[i].TaskID == nil {
unassigned = &groups[i]
} else if *groups[i].TaskID == taskA.ID {
live = &groups[i]
}
}
if live == nil || live.Count != 2 || live.Critical != 1 || live.High != 1 || live.TaskDescription != "group A" {
t.Fatalf("live group mismatch: %+v", live)
}
if unassigned == nil || unassigned.Count != 2 || unassigned.Medium != 1 || unassigned.Low != 1 {
t.Fatalf("unassigned group mismatch: %+v", unassigned)
}
groupForTask := func(items []FindingGroup, taskID int64) *FindingGroup {
for i := range items {
if items[i].TaskID != nil && *items[i].TaskID == taskID {
return &items[i]
}
}
return nil
}
if err := d.SetPaused(taskA.ID, true); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "paused" {
t.Fatalf("paused task group status mismatch: %+v", groups)
}
if err := d.SetPaused(taskA.ID, false); err != nil {
t.Fatal(err)
}
if err := d.Enqueue(taskA.ID, "resume"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "queued" {
t.Fatalf("queued task group status mismatch: %+v", groups)
}
if err := d.SetStatus(taskA.ID, "done"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "done" {
t.Fatalf("terminal task status must win over queued flag: %+v", groups)
}
orphans, total, err := d.ListFindingsPage(FindingFilter{VulnClass: vc, TaskID: FindingUnassignedTask}, 1, 10)
if err != nil || total != 2 || len(orphans) != 2 {
t.Fatalf("unassigned page: len=%d total=%d err=%v", len(orphans), total, err)
}
maxPage := int(^uint(0) >> 1)
farFindings, farTotal, err := d.ListFindingsPage(FindingFilter{VulnClass: vc}, maxPage, 200)
if err != nil || farTotal != len(seed) || len(farFindings) != 0 {
t.Fatalf("far finding page: len=%d total=%d err=%v", len(farFindings), farTotal, err)
}
farGroups, farGroupTotal, farFindingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, maxPage, 100)
if err != nil || farGroupTotal != 2 || farFindingTotal != len(seed) || len(farGroups) != 0 {
t.Fatalf("far group page: len=%d groups=%d findings=%d err=%v",
len(farGroups), farGroupTotal, farFindingTotal, err)
}
filtered, filteredGroups, filteredFindings, err := d.ListFindingGroups(
FindingFilter{VulnClass: vc, Status: FindingResolved}, 1, 10,
)
if err != nil || filteredGroups != 1 || filteredFindings != 1 || len(filtered) != 1 || filtered[0].High != 1 {
t.Fatalf("filtered groups: %+v groups=%d findings=%d err=%v", filtered, filteredGroups, filteredFindings, err)
}
}
func TestAddFindingFollowUpIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
task, err := d.CreateTask("finding follow-up", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("follow-up-%d.example.test", time.Now().UnixNano()),
TaskID: task.ID,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id=$1`, assetID)
store := d.Exploration(task.ExplorationID)
findingNodeID, err := store.AddNode(KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
findingID, err := d.AddFinding(task.ID, findingNodeID, "test", "source", SeverityHigh, "source", "", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
auditInput := Activity{Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "验证可利用性并形成证据链"}
intentID, audit, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "验证可利用性并形成证据链", auditInput)
if err != nil {
t.Fatal(err)
}
secondID, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "从另一条路径深入", Activity{
Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "从另一条路径深入",
})
if err != nil {
t.Fatal(err)
}
if secondID == intentID {
t.Fatal("repeated follow-up submissions must create distinct intents")
}
expectedAuditSummary := fmt.Sprintf("%s #%d", auditInput.Summary, intentID)
if audit.ID <= 0 || audit.NodeID == nil || *audit.NodeID != intentID || audit.CreatedAt.IsZero() || audit.Summary != expectedAuditSummary {
t.Fatalf("persisted audit mismatch: %+v", audit)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil {
t.Fatalf("GetNode: node=%+v err=%v", node, err)
}
if node.Kind != KindIntent || node.Priority != 10 || node.State != "open" || node.Origin != "human" {
t.Fatalf("follow-up intent metadata: %+v", node)
}
var anchorCount, edgeCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1 AND asset_id=$2`, intentID, assetID).Scan(&anchorCount); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_edges
WHERE exploration_id=$1 AND src_id=$2 AND rel=$3 AND dst_id=$4`,
task.ExplorationID, findingNodeID, RelDerivedFrom, intentID).Scan(&edgeCount); err != nil {
t.Fatal(err)
}
if anchorCount != 1 || edgeCount != 1 {
t.Fatalf("follow-up lineage: anchors=%d edges=%d", anchorCount, edgeCount)
}
var persistedAuditCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM activity
WHERE id=$1 AND exploration_id=$2 AND node_id=$3 AND worker='system' AND summary=$4`,
audit.ID, task.ExplorationID, intentID, expectedAuditSummary).Scan(&persistedAuditCount); err != nil {
t.Fatal(err)
}
if persistedAuditCount != 1 {
t.Fatalf("atomic audit count=%d, want 1", persistedAuditCount)
}
var nodesBefore, activitiesBefore int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesBefore); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesBefore); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "must roll back", Activity{
Worker: "system", Kind: "text", Summary: "invalid audit", Metadata: json.RawMessage(`{`),
}); err == nil {
t.Fatal("invalid activity metadata should fail the transaction")
}
var nodesAfter, activitiesAfter int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesAfter); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesAfter); err != nil {
t.Fatal(err)
}
if nodesAfter != nodesBefore || activitiesAfter != activitiesBefore {
t.Fatalf("activity failure did not roll back graph: nodes %d->%d activities %d->%d",
nodesBefore, nodesAfter, activitiesBefore, activitiesAfter)
}
if _, err := d.Exec(`DELETE FROM exploration_nodes WHERE id=$1`, findingNodeID); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "should fail", auditInput); !errors.Is(err, ErrFindingOriginUnavailable) {
t.Fatalf("missing source node: want ErrFindingOriginUnavailable, got %v", err)
}
}
+32
View File
@@ -0,0 +1,32 @@
package db
import "fmt"
// DiscardOpenIntent compensates a follow-up creation when task admission fails.
// The task execution gate must still be held by the caller, so the intent cannot
// be claimed between this check and deletion. Edges and anchors cascade with the
// node; activity is deleted explicitly because its node FK otherwise becomes NULL.
func (s *ExplorationStore) DiscardOpenIntent(id int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil {
return err
}
res, err := tx.Exec(`DELETE FROM exploration_nodes
WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID)
if err != nil {
return err
}
removed, err := res.RowsAffected()
if err != nil {
return err
}
if removed != 1 {
return fmt.Errorf("open intent %d was not available for admission rollback", id)
}
return tx.Commit()
}
+285
View File
@@ -0,0 +1,285 @@
package db
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("intent CAS", "only one controller wins")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
const contenders = 12
var winners atomic.Int32
var wg sync.WaitGroup
start := make(chan struct{})
for range contenders {
wg.Add(1)
go func() {
defer wg.Done()
<-start
changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
if transitionErr != nil {
t.Errorf("transition: %v", transitionErr)
return
}
if changed {
winners.Add(1)
}
}()
}
close(start)
wg.Wait()
if got := winners.Load(); got != 1 {
t.Fatalf("CAS winners=%d, want 1", got)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil || node.State != "paused" {
t.Fatalf("node=%+v err=%v, want paused", node, err)
}
}
func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
baselineDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
baseline := dailyTokenBuckets(baselineDaily)
appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
t.Helper()
activityID, appendErr := store.AppendActivity(Activity{
NodeID: &intentID, Worker: "work#1", Kind: kind,
InputTokens: input, OutputTokens: output,
CacheReadTokens: read, CacheWriteTokens: write,
})
if appendErr != nil {
t.Fatal(appendErr)
}
if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
t.Fatal(updateErr)
}
}
values := func(input, output, read, write int) (*int, *int, *int, *int) {
return &input, &output, &read, &write
}
now := time.Now().UTC()
atUTCNoon := func(daysAgo int) time.Time {
at := now.AddDate(0, 0, -daysAgo)
return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
}
authoritativeDay := atUTCNoon(6)
fallbackDay := atUTCNoon(4)
trailingDay := atUTCNoon(2)
// Run 1 has an authoritative result; its preceding cumulative frame must not
// be counted a second time or move usage to the frame's earlier date.
i, o, r, w := values(100, 100, 100, 100)
appendUsage("usage", atUTCNoon(7), i, o, r, w)
i, o, r, w = values(10, 11, 12, 13)
appendUsage("result", authoritativeDay, i, o, r, w)
// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
// to the result's date, not the cumulative frame's date.
i, o, r, w = values(20, 21, 22, 23)
appendUsage("usage", atUTCNoon(5), i, o, r, w)
appendUsage("result", fallbackDay, nil, nil, nil, nil)
// Run 3 was interrupted before a terminal result was persisted.
i, o, r, w = values(30, 31, 32, 33)
appendUsage("usage", trailingDay, i, o, r, w)
beforeCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
beforeCancel := dailyTokenBuckets(beforeCancelDaily)
assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
if err := store.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(intentID)
if err != nil {
t.Fatal(err)
}
if cleanup.Activities != 5 {
t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total, 60, 63, 66, 69)
afterCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
afterCancel := dailyTokenBuckets(afterCancelDaily)
assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
if len(sessions) != 0 {
t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
}
var rollups, datedRollups int
if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
) FROM activity
WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
t.Fatal(err)
}
if rollups != 3 || datedRollups != rollups {
t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
}
if _, err := store.CancelIntent(intentID); err == nil {
t.Fatal("second cancellation unexpectedly succeeded")
}
afterRetry, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
afterRetryDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
}
func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
out := make(map[string]DailyTokenBucket, len(items))
for _, item := range items {
out[item.Day] = item
}
return out
}
func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
t.Helper()
day := at.UTC().Format(time.DateOnly)
actual, initial := got[day], baseline[day]
if actual.InputTokens-initial.InputTokens != input ||
actual.OutputTokens-initial.OutputTokens != output ||
actual.CacheReadTokens-initial.CacheReadTokens != read {
t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
}
}
func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
}
for day, expected := range want {
if actual, ok := got[day]; !ok || actual != expected {
t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
}
}
}
func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
for _, intentID := range []int64{first, second} {
if err := store.Link(intentID, RelYields, fact); err != nil {
t.Fatal(err)
}
if err := store.Link(intentID, RelYields, finding); err != nil {
t.Fatal(err)
}
}
if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
if err := store.SetIntentState(first, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(first)
if err != nil {
t.Fatal(err)
}
if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
}
for _, nodeID := range []int64{fact, finding} {
node, getErr := store.GetNode(nodeID)
if getErr != nil || node == nil {
t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
}
}
}
+59 -2
View File
@@ -31,6 +31,11 @@ CREATE TABLE IF NOT EXISTS assets (
'root_domain','ip','subdomain','app','service','endpoint'
)),
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
-- explicit: caller/user selected the company; scope: derived from company_scope.
-- Existing installations are conservatively migrated as explicit so a scope
-- rebuild can never erase a historical manual association.
company_source TEXT NOT NULL DEFAULT 'explicit'
CHECK (company_source IN ('explicit','scope')),
task_ids BIGINT[] NOT NULL DEFAULT '{}',
domain TEXT,
root_domain TEXT,
@@ -84,6 +89,13 @@ CREATE INDEX IF NOT EXISTS idx_av2_bound_domains ON assets USING GIN(bound_domai
CREATE INDEX IF NOT EXISTS idx_av2_open_ports ON assets USING GIN(open_ports) WHERE type = 'ip';
CREATE INDEX IF NOT EXISTS idx_av2_last_seen ON assets(last_seen DESC);
CREATE INDEX IF NOT EXISTS idx_av2_type_seen ON assets(type, last_seen DESC);
ALTER TABLE assets ADD COLUMN IF NOT EXISTS company_source TEXT;
UPDATE assets SET company_source = 'explicit' WHERE company_source IS NULL;
ALTER TABLE assets ALTER COLUMN company_source SET DEFAULT 'explicit';
ALTER TABLE assets ALTER COLUMN company_source SET NOT NULL;
ALTER TABLE assets DROP CONSTRAINT IF EXISTS assets_company_source_check;
ALTER TABLE assets ADD CONSTRAINT assets_company_source_check
CHECK (company_source IN ('explicit','scope'));
DROP TRIGGER IF EXISTS trg_av2_upd ON assets;
CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
@@ -91,17 +103,36 @@ CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
CREATE TABLE IF NOT EXISTS company_scope (
id BIGSERIAL PRIMARY KEY,
company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr')),
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr','icp','keyword')),
domain TEXT,
net CIDR,
value TEXT,
raw TEXT NOT NULL,
reason TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT uq_sv2_domain UNIQUE (company_id, domain),
CONSTRAINT uq_sv2_net UNIQUE (company_id, net)
CONSTRAINT uq_sv2_net UNIQUE (company_id, net),
CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
)
);
-- Existing installations need the new text payload and expanded kind check.
ALTER TABLE company_scope ADD COLUMN IF NOT EXISTS value TEXT;
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS company_scope_kind_check;
ALTER TABLE company_scope ADD CONSTRAINT company_scope_kind_check
CHECK (kind IN ('domain','ip','cidr','icp','keyword'));
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS ck_company_scope_payload;
ALTER TABLE company_scope ADD CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
);
CREATE INDEX IF NOT EXISTS idx_sv2_domain ON company_scope(domain) WHERE kind = 'domain';
CREATE INDEX IF NOT EXISTS idx_sv2_net ON company_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr');
CREATE UNIQUE INDEX IF NOT EXISTS uq_sv2_value ON company_scope(company_id, kind, value) WHERE kind IN ('icp','keyword');
CREATE INDEX IF NOT EXISTS idx_sv2_icp ON company_scope(value) WHERE kind = 'icp';
CREATE INDEX IF NOT EXISTS idx_sv2_company ON company_scope(company_id);
-- =====================================================================
@@ -301,6 +332,8 @@ CREATE TABLE IF NOT EXISTS tasks (
CHECK (status IN ('created','running','paused','done','failed','timeout')),
paused BOOLEAN NOT NULL DEFAULT false,
queued BOOLEAN NOT NULL DEFAULT false,
queued_at TIMESTAMPTZ,
queue_mode TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
llm_chain_revision BIGINT NOT NULL DEFAULT 0,
@@ -324,9 +357,30 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300;
-- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false;
-- 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 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 (
@@ -573,10 +627,13 @@ CREATE TABLE IF NOT EXISTS conversations (
agent_key TEXT NOT NULL,
title TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
pinned_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE conversations ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ;
CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned_at DESC) WHERE pinned_at IS NOT NULL;
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
+33 -2
View File
@@ -60,6 +60,11 @@ func (d *DB) hydrateTaskContext(t *Task) error {
return err
}
t.SourceTaskIDs = sources
companies, err := d.TaskCompanyIDs(t.ID)
if err != nil {
return err
}
t.CompanyIDs = companies
t.LLMProfileIDs = make([]int64, 0, len(chain))
t.LLMFailoverState = "default"
var latest *TaskLLMProfile
@@ -172,6 +177,28 @@ func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) {
return out, rows.Err()
}
// TaskCompanyIDs returns the companies whose asset scopes are available to the
// task. The task_scope rows remain the single source of truth.
func (d *DB) TaskCompanyIDs(taskID int64) ([]int64, error) {
rows, err := d.Query(`
SELECT company_id FROM task_scope
WHERE task_id=$1 AND kind='company' AND company_id IS NOT NULL
ORDER BY id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) {
rows, err := d.Query(`
SELECT t.id, t.exploration_id, t.description, t.goal, t.status
@@ -339,13 +366,17 @@ ORDER BY position LIMIT 1`, taskID, position).Scan(&next)
case nil:
out.Advanced = true
out.NextProfileID = &next
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=$2, llm_profile_id=$2 WHERE id=$1`, taskID, next); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID, next); err != nil {
return out, err
}
case sql.ErrNoRows:
out.Advanced = true
out.ChainExhausted = true
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=NULL, llm_profile_id=NULL WHERE id=$1`, taskID); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID); err != nil {
return out, err
}
default:
+66
View File
@@ -0,0 +1,66 @@
package db
import (
"testing"
"time"
)
func TestTaskQueuePreservesBootstrapAndFIFOPosition(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("queue metadata", "keep first-run mode", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
first, err := d.GetTask(task.ID)
if err != nil || first == nil || first.QueuedAt == nil {
t.Fatalf("first enqueue: task=%+v err=%v", first, err)
}
// A follow-up/rerun can try to admit an already queued task as resume. The
// original bootstrap mode and FIFO timestamp must remain authoritative.
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "resume"); err != nil {
t.Fatal(err)
}
second, err := d.GetTask(task.ID)
if err != nil || second == nil {
t.Fatalf("second enqueue: task=%+v err=%v", second, err)
}
if second.QueueMode != "bootstrap" {
t.Fatalf("queue mode=%q, want bootstrap", second.QueueMode)
}
if second.QueuedAt == nil || !second.QueuedAt.Equal(*first.QueuedAt) {
t.Fatalf("repeated enqueue moved FIFO position: first=%v second=%v", first.QueuedAt, second.QueuedAt)
}
// Pausing a queued task removes it from the queue but keeps its required
// startup mode. A later requeue receives a new tail position.
if err := d.Dequeue(task.ID, false); err != nil {
t.Fatal(err)
}
paused, err := d.GetTask(task.ID)
if err != nil || paused == nil || paused.Queued || paused.QueueMode != "bootstrap" {
t.Fatalf("paused queue metadata: task=%+v err=%v", paused, err)
}
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
requeued, err := d.GetTask(task.ID)
if err != nil || requeued == nil || requeued.QueuedAt == nil {
t.Fatalf("requeue: task=%+v err=%v", requeued, err)
}
if !requeued.QueuedAt.After(*first.QueuedAt) {
t.Fatalf("requeue did not move to FIFO tail: first=%v requeued=%v", first.QueuedAt, requeued.QueuedAt)
}
}
+27
View File
@@ -1,6 +1,8 @@
package db
import (
"errors"
"reflect"
"strings"
"testing"
)
@@ -18,3 +20,28 @@ func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) {
t.Fatalf("expected source-count validation error, got %v", err)
}
}
func TestNormalizeTaskCompanyIDs(t *testing.T) {
got, err := NormalizeTaskCompanyIDs([]int64{4, 2, 4, 7, 2})
if err != nil {
t.Fatal(err)
}
if want := []int64{4, 2, 7}; !reflect.DeepEqual(got, want) {
t.Fatalf("NormalizeTaskCompanyIDs=%v, want %v", got, want)
}
if _, err := NormalizeTaskCompanyIDs([]int64{1, 0}); !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("invalid company id error=%v", err)
}
}
func TestCreateTaskRejectsTooManyCompaniesBeforeOpeningTransaction(t *testing.T) {
companyIDs := make([]int64, MaxTaskCompanyCount+1)
for i := range companyIDs {
companyIDs[i] = int64(i + 1)
}
_, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{CompanyIDs: companyIDs})
if !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("expected company-count validation error, got %v", err)
}
}
+231
View File
@@ -0,0 +1,231 @@
package db
import (
"database/sql"
"errors"
"fmt"
"strings"
"time"
"unicode/utf8"
"github.com/jackc/pgx/v5/pgconn"
)
const (
MaxTaskTemplateNameRunes = 120
MaxTaskTemplateTextRunes = 16000
)
var (
ErrTaskTemplateInvalid = errors.New("invalid task template")
ErrTaskTemplateNameConflict = errors.New("task template name already exists")
ErrTaskTemplateNotFound = errors.New("task template not found")
)
// TaskTemplate is a reusable task description and goal preset.
type TaskTemplate struct {
ID int64 `json:"id"`
Name string `json:"name"`
NKey string `json:"-"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// TaskTemplateInput is the create/update payload after normalization.
type TaskTemplateInput struct {
Name string
Description string
Goal string
}
// TaskTemplatePatch changes only fields whose pointers are non-nil.
type TaskTemplatePatch struct {
Name *string
Description *string
Goal *string
}
const taskTemplateCols = `id, name, nkey, description, goal, created_at, updated_at`
func scanTaskTemplate(row interface{ Scan(...any) error }) (TaskTemplate, error) {
var t TaskTemplate
err := row.Scan(&t.ID, &t.Name, &t.NKey, &t.Description, &t.Goal, &t.CreatedAt, &t.UpdatedAt)
return t, err
}
// taskTemplateName normalizes display whitespace while preserving the user's case.
func taskTemplateName(name string) string { return strings.Join(strings.Fields(name), " ") }
// taskTemplateNKey is the case-insensitive identity used by the unique index.
func taskTemplateNKey(name string) string { return strings.ToLower(taskTemplateName(name)) }
func normalizeTaskTemplateInput(in TaskTemplateInput) (TaskTemplateInput, string, error) {
in.Name = taskTemplateName(in.Name)
in.Description = strings.TrimSpace(in.Description)
in.Goal = strings.TrimSpace(in.Goal)
switch {
case in.Name == "":
return in, "", fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Name) > MaxTaskTemplateNameRunes:
return in, "", fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
case in.Description == "":
return in, "", fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Description) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
case in.Goal == "":
return in, "", fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Goal) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
return in, taskTemplateNKey(in.Name), nil
}
func normalizeTaskTemplatePatch(patch TaskTemplatePatch) (TaskTemplatePatch, *string, error) {
if patch.Name == nil && patch.Description == nil && patch.Goal == nil {
return patch, nil, fmt.Errorf("%w: no fields supplied", ErrTaskTemplateInvalid)
}
var nkey *string
if patch.Name != nil {
name := taskTemplateName(*patch.Name)
if name == "" {
return patch, nil, fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(name) > MaxTaskTemplateNameRunes {
return patch, nil, fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
}
key := taskTemplateNKey(name)
patch.Name = &name
nkey = &key
}
if patch.Description != nil {
description := strings.TrimSpace(*patch.Description)
if description == "" {
return patch, nil, fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(description) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Description = &description
}
if patch.Goal != nil {
goal := strings.TrimSpace(*patch.Goal)
if goal == "" {
return patch, nil, fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(goal) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Goal = &goal
}
return patch, nkey, nil
}
func taskTemplateUniqueViolation(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "23505"
}
// CreateTaskTemplate inserts one globally reusable preset.
func (d *DB) CreateTaskTemplate(in TaskTemplateInput) (*TaskTemplate, error) {
in, nkey, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`
INSERT INTO task_templates(name, nkey, description, goal)
VALUES ($1,$2,$3,$4)
ON CONFLICT (nkey) DO NOTHING
RETURNING `+taskTemplateCols, in.Name, nkey, in.Description, in.Goal))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// ListTaskTemplates returns the most recently maintained templates first.
func (d *DB) ListTaskTemplates() ([]*TaskTemplate, error) {
rows, err := d.Query(`SELECT ` + taskTemplateCols + ` FROM task_templates ORDER BY updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*TaskTemplate{}
for rows.Next() {
t, err := scanTaskTemplate(rows)
if err != nil {
return nil, err
}
out = append(out, &t)
}
return out, rows.Err()
}
// GetTaskTemplate returns nil when id does not exist.
func (d *DB) GetTaskTemplate(id int64) (*TaskTemplate, error) {
t, err := scanTaskTemplate(d.QueryRow(`SELECT `+taskTemplateCols+` FROM task_templates WHERE id=$1`, id))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &t, nil
}
// UpdateTaskTemplate replaces the editable fields of one preset.
func (d *DB) UpdateTaskTemplate(id int64, in TaskTemplateInput) (*TaskTemplate, error) {
in, _, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
return d.PatchTaskTemplate(id, TaskTemplatePatch{
Name: &in.Name, Description: &in.Description, Goal: &in.Goal,
})
}
// PatchTaskTemplate atomically changes only the supplied fields. Keeping the
// merge in one UPDATE prevents concurrent disjoint PATCH requests from losing
// each other's changes.
func (d *DB) PatchTaskTemplate(id int64, patch TaskTemplatePatch) (*TaskTemplate, error) {
patch, nkey, err := normalizeTaskTemplatePatch(patch)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`UPDATE task_templates
SET name=CASE WHEN $2 THEN $3::text ELSE name END,
nkey=CASE WHEN $2 THEN $4::text ELSE nkey END,
description=CASE WHEN $5 THEN $6::text ELSE description END,
goal=CASE WHEN $7 THEN $8::text ELSE goal END
WHERE id=$1
RETURNING `+taskTemplateCols,
id,
patch.Name != nil, patch.Name, nkey,
patch.Description != nil, patch.Description,
patch.Goal != nil, patch.Goal,
))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNotFound
}
if taskTemplateUniqueViolation(err) {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// DeleteTaskTemplate deletes one preset and reports whether it existed.
func (d *DB) DeleteTaskTemplate(id int64) (bool, error) {
result, err := d.Exec(`DELETE FROM task_templates WHERE id=$1`, id)
if err != nil {
return false, err
}
n, err := result.RowsAffected()
return n > 0, err
}
+131
View File
@@ -0,0 +1,131 @@
package db
import (
"errors"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestTaskTemplateCRUDAndNormalizedUniqueness(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
name := fmt.Sprintf("Template %d", suffix)
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: " " + strings.ReplaceAll(name, " ", " ") + " ",
Description: " initial description ",
Goal: " initial goal ",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
if created.Name != name || created.Description != "initial description" || created.Goal != "initial goal" {
t.Fatalf("template was not normalized: %+v", created)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: strings.ToUpper(name), Description: "duplicate", Goal: "duplicate",
}); !errors.Is(err, ErrTaskTemplateNameConflict) {
t.Fatalf("duplicate create error = %v, want %v", err, ErrTaskTemplateNameConflict)
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil || got == nil || got.Name != name {
t.Fatalf("GetTaskTemplate = %+v, %v", got, err)
}
listed, err := d.ListTaskTemplates()
if err != nil {
t.Fatal(err)
}
found := false
for _, template := range listed {
if template.ID == created.ID {
found = true
break
}
}
if !found {
t.Fatalf("created template %d missing from list", created.ID)
}
updatedName := name + " updated"
updated, err := d.UpdateTaskTemplate(created.ID, TaskTemplateInput{
Name: updatedName, Description: "new description", Goal: "new goal",
})
if err != nil {
t.Fatal(err)
}
if updated.Name != updatedName || updated.Description != "new description" || updated.Goal != "new goal" {
t.Fatalf("unexpected updated template: %+v", updated)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{Name: "", Description: "x", Goal: "y"}); !errors.Is(err, ErrTaskTemplateInvalid) {
t.Fatalf("empty name error = %v, want %v", err, ErrTaskTemplateInvalid)
}
deleted, err := d.DeleteTaskTemplate(created.ID)
if err != nil || !deleted {
t.Fatalf("DeleteTaskTemplate = %v, %v", deleted, err)
}
deleted, err = d.DeleteTaskTemplate(created.ID)
if err != nil || deleted {
t.Fatalf("second DeleteTaskTemplate = %v, %v", deleted, err)
}
}
func TestTaskTemplateDisjointPatchesCompose(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: fmt.Sprintf("Concurrent template %d", time.Now().UnixNano()),
Description: "initial description",
Goal: "initial goal",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
description := "description from concurrent patch"
goal := "goal from concurrent patch"
start := make(chan struct{})
errs := make(chan error, 2)
var wg sync.WaitGroup
for _, patch := range []TaskTemplatePatch{{Description: &description}, {Goal: &goal}} {
patch := patch
wg.Add(1)
go func() {
defer wg.Done()
<-start
_, err := d.PatchTaskTemplate(created.ID, patch)
errs <- err
}()
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatal(err)
}
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil {
t.Fatal(err)
}
if got == nil || got.Description != description || got.Goal != goal {
t.Fatalf("disjoint patches lost an update: %+v", got)
}
}
+119 -12
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"time"
@@ -10,14 +11,16 @@ import (
// Task is a row in the task registry (1:1 with an exploration).
type Task struct {
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
QueuedAt *time.Time `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
// for ActiveLLMProfileID while older API clients still send one profile id.
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
@@ -26,6 +29,7 @@ type Task struct {
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
@@ -74,6 +78,40 @@ const MinPlanHeartbeatSeconds = 600
// multiplying graph and asset-context queries without limit.
const MaxTaskSourceCount = 8
// MaxTaskCompanyCount bounds the number of company asset scopes attached to a
// task. Company scopes are prompt context, so an unbounded list would inflate
// every planner and worker call.
const MaxTaskCompanyCount = 32
var (
ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids")
ErrTaskCompanyNotFound = errors.New("task company not found")
)
// NormalizeTaskCompanyIDs validates IDs and removes duplicates while retaining
// the user's first-seen order.
func NormalizeTaskCompanyIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, nil
}
seen := make(map[int64]struct{}, min(len(ids), MaxTaskCompanyCount))
normalized := make([]int64, 0, min(len(ids), MaxTaskCompanyCount))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("%w: company id must be positive", ErrTaskCompanyIDsInvalid)
}
if _, exists := seen[id]; exists {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
if len(normalized) > MaxTaskCompanyCount {
return nil, fmt.Errorf("%w: got more than %d unique companies", ErrTaskCompanyIDsInvalid, MaxTaskCompanyCount)
}
}
return normalized, nil
}
func normalizeHeartbeat(sec int) int {
if sec < MinPlanHeartbeatSeconds {
return MinPlanHeartbeatSeconds
@@ -95,6 +133,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe
// with the task/exploration row.
type TaskCreateOptions struct {
SourceTaskIDs []int64
CompanyIDs []int64
LLMProfileIDs []int64
TimeoutSeconds int
PlanHeartbeatSeconds int
@@ -106,6 +145,11 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti
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
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -144,6 +188,7 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload));
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,
}
if err := tx.QueryRow(`
@@ -155,6 +200,9 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil {
return nil, err
}
if err := insertTaskCompanies(tx, t.ID, opts.CompanyIDs); err != nil {
return nil, err
}
if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil {
return nil, err
}
@@ -166,6 +214,33 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
return t, tx.Commit()
}
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 {
@@ -199,11 +274,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
return nil
}
const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
const taskCols = `id, description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
return nil, err
}
return &t, nil
@@ -263,11 +338,43 @@ func (d *DB) SetPaused(id int64, paused bool) error {
return err
}
func (d *DB) SetQueued(id int64, queued bool) error {
_, err := d.Exec(`UPDATE tasks SET queued=$1 WHERE id=$2`, queued, id)
// Enqueue places a task at the tail of the persistent admission queue. Repeating
// the operation while it is already queued keeps its original FIFO position.
func (d *DB) Enqueue(id int64, mode string) error {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
_, err := d.Exec(`UPDATE tasks
SET queued=true,
queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END,
queue_mode=CASE
WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END
WHERE id=$1`, id, mode)
return err
}
// Dequeue removes the concurrency hold. clearMode is false when a user pauses a
// queued task, preserving whether its next admission must bootstrap or resume.
func (d *DB) Dequeue(id int64, clearMode bool) error {
_, err := d.Exec(`UPDATE tasks
SET queued=false,
queued_at=NULL,
queue_mode=CASE WHEN $2 THEN '' ELSE queue_mode END
WHERE id=$1`, id, clearMode)
return err
}
// SetQueued is the compatibility helper used by older callers and tests. New
// scheduling code should use Enqueue/Dequeue so FIFO metadata is explicit.
func (d *DB) SetQueued(id int64, queued bool) error {
if queued {
return d.Enqueue(id, "bootstrap")
}
return d.Dequeue(id, true)
}
// SetStatus updates a task's lifecycle status. Entering a terminal state
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
+80
View File
@@ -1,6 +1,7 @@
package db
import (
"errors"
"fmt"
"sync"
"testing"
@@ -277,6 +278,7 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] {
t.Fatalf("replacement revision/cursor not preserved: %+v", got)
}
replacementRevision := got.LLMChainRevision
transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota")
if err != nil {
@@ -285,6 +287,13 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] {
t.Fatalf("unexpected first transition: %+v", transition)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.LLMChainRevision != replacementRevision+1 {
t.Fatalf("automatic cursor advance did not increment revision: before=%d after=%d", replacementRevision, got.LLMChainRevision)
}
late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate")
if err != nil {
t.Fatal(err)
@@ -483,3 +492,74 @@ func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) {
t.Fatal("terminal task LLM edit should be rejected")
}
}
func TestCreateTaskWithCompanyScopes(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "")
if err != nil {
t.Fatal(err)
}
companyB, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company B %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB)
})
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyB},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("creation result company IDs=%v", task.CompanyIDs)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("hydrated company IDs=%v", got.CompanyIDs)
}
var scopeCount int
if err := d.QueryRow(`SELECT count(*) FROM task_scope WHERE task_id=$1 AND kind='company'`, task.ID).Scan(&scopeCount); err != nil {
t.Fatal(err)
}
if scopeCount != 2 {
t.Fatalf("company task scope rows=%d want 2", scopeCount)
}
duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyA, companyB, companyA},
})
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(duplicateTask.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("company IDs were not normalized: %v", duplicateTask.CompanyIDs)
}
if err := d.DeleteTask(duplicateTask.ID); err != nil {
t.Fatal(err)
}
badDescription := fmt.Sprintf("invalid-company-%d", suffix)
if _, err := d.CreateTaskWithOptions(badDescription, "rollback", TaskCreateOptions{CompanyIDs: []int64{companyA, 1 << 62}}); !errors.Is(err, ErrTaskCompanyNotFound) {
t.Fatalf("missing company error=%v, want %v", err, ErrTaskCompanyNotFound)
}
var leaked int
if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, badDescription).Scan(&leaked); err != nil {
t.Fatal(err)
}
if leaked != 0 {
t.Fatalf("failed company association leaked %d exploration rows", leaked)
}
}