feat(tasks): add templates findings workflows and batch controls
This commit is contained in:
1 parent
b7e0b93023
commit
d3183dbbd2
65 files changed
+11603
-2198
No files matched your search
+16
-1
@@ -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
@@ -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))
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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 导出(导出当前筛选/全部)。结果按严重等级降序、
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package db
|
||||
|
||||
import "fmt"
|
||||
|
||||
// DiscardOpenIntent compensates a follow-up creation when task admission fails.
|
||||
// The task execution gate must still be held by the caller, so the intent cannot
|
||||
// be claimed between this check and deletion. Edges and anchors cascade with the
|
||||
// node; activity is deleted explicitly because its node FK otherwise becomes NULL.
|
||||
func (s *ExplorationStore) DiscardOpenIntent(id int64) error {
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := tx.Exec(`DELETE FROM exploration_nodes
|
||||
WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
removed, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if removed != 1 {
|
||||
return fmt.Errorf("open intent %d was not available for admission rollback", id)
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
expID, err := d.CreateExploration("intent CAS", "only one controller wins")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
|
||||
store := d.Exploration(expID)
|
||||
intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
|
||||
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
|
||||
}
|
||||
|
||||
const contenders = 12
|
||||
var winners atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
for range contenders {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
|
||||
if transitionErr != nil {
|
||||
t.Errorf("transition: %v", transitionErr)
|
||||
return
|
||||
}
|
||||
if changed {
|
||||
winners.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
if got := winners.Load(); got != 1 {
|
||||
t.Fatalf("CAS winners=%d, want 1", got)
|
||||
}
|
||||
node, err := store.GetNode(intentID)
|
||||
if err != nil || node == nil || node.State != "paused" {
|
||||
t.Fatalf("node=%+v err=%v, want paused", node, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
|
||||
store := d.Exploration(expID)
|
||||
intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
|
||||
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
|
||||
}
|
||||
|
||||
baselineDaily, err := d.TokenDailyAll(30)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
baseline := dailyTokenBuckets(baselineDaily)
|
||||
|
||||
appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
|
||||
t.Helper()
|
||||
activityID, appendErr := store.AppendActivity(Activity{
|
||||
NodeID: &intentID, Worker: "work#1", Kind: kind,
|
||||
InputTokens: input, OutputTokens: output,
|
||||
CacheReadTokens: read, CacheWriteTokens: write,
|
||||
})
|
||||
if appendErr != nil {
|
||||
t.Fatal(appendErr)
|
||||
}
|
||||
if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
|
||||
t.Fatal(updateErr)
|
||||
}
|
||||
}
|
||||
values := func(input, output, read, write int) (*int, *int, *int, *int) {
|
||||
return &input, &output, &read, &write
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
atUTCNoon := func(daysAgo int) time.Time {
|
||||
at := now.AddDate(0, 0, -daysAgo)
|
||||
return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
|
||||
}
|
||||
authoritativeDay := atUTCNoon(6)
|
||||
fallbackDay := atUTCNoon(4)
|
||||
trailingDay := atUTCNoon(2)
|
||||
|
||||
// Run 1 has an authoritative result; its preceding cumulative frame must not
|
||||
// be counted a second time or move usage to the frame's earlier date.
|
||||
i, o, r, w := values(100, 100, 100, 100)
|
||||
appendUsage("usage", atUTCNoon(7), i, o, r, w)
|
||||
i, o, r, w = values(10, 11, 12, 13)
|
||||
appendUsage("result", authoritativeDay, i, o, r, w)
|
||||
// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
|
||||
// to the result's date, not the cumulative frame's date.
|
||||
i, o, r, w = values(20, 21, 22, 23)
|
||||
appendUsage("usage", atUTCNoon(5), i, o, r, w)
|
||||
appendUsage("result", fallbackDay, nil, nil, nil, nil)
|
||||
// Run 3 was interrupted before a terminal result was persisted.
|
||||
i, o, r, w = values(30, 31, 32, 33)
|
||||
appendUsage("usage", trailingDay, i, o, r, w)
|
||||
|
||||
beforeCancelDaily, err := d.TokenDailyAll(30)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
beforeCancel := dailyTokenBuckets(beforeCancelDaily)
|
||||
assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
|
||||
assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
|
||||
assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
|
||||
|
||||
if err := store.SetIntentState(intentID, "paused"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cleanup, err := store.CancelIntent(intentID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cleanup.Activities != 5 {
|
||||
t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
|
||||
}
|
||||
total, err := store.TokenTotal()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertTokenUsage(t, total, 60, 63, 66, 69)
|
||||
afterCancelDaily, err := d.TokenDailyAll(30)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
afterCancel := dailyTokenBuckets(afterCancelDaily)
|
||||
assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
|
||||
assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
|
||||
assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
|
||||
assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
|
||||
|
||||
sessions, err := store.TokenStatsBySession()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sessions) != 0 {
|
||||
t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
|
||||
}
|
||||
var rollups, datedRollups int
|
||||
if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
|
||||
WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
|
||||
) FROM activity
|
||||
WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
|
||||
AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rollups != 3 || datedRollups != rollups {
|
||||
t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
|
||||
}
|
||||
if _, err := store.CancelIntent(intentID); err == nil {
|
||||
t.Fatal("second cancellation unexpectedly succeeded")
|
||||
}
|
||||
afterRetry, err := store.TokenTotal()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
|
||||
afterRetryDaily, err := d.TokenDailyAll(30)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
|
||||
}
|
||||
|
||||
func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
|
||||
out := make(map[string]DailyTokenBucket, len(items))
|
||||
for _, item := range items {
|
||||
out[item.Day] = item
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
|
||||
t.Helper()
|
||||
day := at.UTC().Format(time.DateOnly)
|
||||
actual, initial := got[day], baseline[day]
|
||||
if actual.InputTokens-initial.InputTokens != input ||
|
||||
actual.OutputTokens-initial.OutputTokens != output ||
|
||||
actual.CacheReadTokens-initial.CacheReadTokens != read {
|
||||
t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
|
||||
day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
|
||||
actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
|
||||
}
|
||||
}
|
||||
|
||||
func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
|
||||
t.Helper()
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
|
||||
}
|
||||
for day, expected := range want {
|
||||
if actual, ok := got[day]; !ok || actual != expected {
|
||||
t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
|
||||
store := d.Exploration(expID)
|
||||
first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, intentID := range []int64{first, second} {
|
||||
if err := store.Link(intentID, RelYields, fact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.Link(intentID, RelYields, finding); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
|
||||
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
|
||||
}
|
||||
if err := store.SetIntentState(first, "paused"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cleanup, err := store.CancelIntent(first)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
|
||||
t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
|
||||
}
|
||||
for _, nodeID := range []int64{fact, finding} {
|
||||
node, getErr := store.GetNode(nodeID)
|
||||
if getErr != nil || node == nil {
|
||||
t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
+59
-2
@@ -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
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user