admin/ARTEX-chixuweihu
1
0 Fork 0
mirror of https://github.com/RuoJi6/ARTEX.git synced 2026-10-28 01:03:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
pphua committed 2026-08-21 21:01:39 +08:00
1 parent b7e0b93023
commit d3183dbbd2
65 files changed
+11603 -2198
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CHANGELOG.md
+25
View File
@@ -2,6 +2,31 @@
本项目的重要变更记录在此文件中，格式参考 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/)。
## [Unreleased]
### Added
- 增加全局任务模板管理，新建任务可直接加载预设描述与目标，也可将当前内容另存为模板。
- 新建任务支持关联多个企业资产范围；企业范围采用结构化规则，覆盖域名、IP、CIDR、ICP 和企业关键词。
- 发现页按任务分组展示漏洞并支持组内独立分页；漏洞可提交带描述的高优先级深入意图，在原任务中继续验证。
- 对话支持重命名、置顶和取消置顶；任务与 Worker 会话支持批量暂停、继续及逐项结果反馈。
- 会话列表展示 Main Agent、Planner 和 Worker 下一次调用的模型，并提供完整历史 Token 聚合与高亮统计。
### Changed
- 任务列表移除单独的“进入”操作，改为通过可聚焦的任务标题链接进入详情。
- 新增企业改用与任务、LLM 一致的右侧抽屉，并移除 Logo URL 输入项。
- 任务创建、恢复、漏洞深入和队列补位统一使用持久化并发准入与 FIFO 调度路径。
- 任务资产、企业资产、发现任务组及组内漏洞使用服务端真实分页、稳定排序和精确总数。
### Fixed
- 修复 LLM 配置链耗尽错误在窄屏横向溢出，以及会话、对话头部和企业范围编辑在移动端的同类布局问题。
- 修复跨任务同名漏洞节点可能联动编辑、删除或展开，以及删除末页数据后页码越界的问题。
- 修复任务生命周期并发写入可能以旧快照覆盖终态、LLM 自动切换可能被旧配置快照回退的问题。
- 修复 `timeout` 任务复活后仍被旧截止时间和收尾状态阻塞；复活后会保留时限配置并在下一次真实执行时重新计时。
- 修复资产分页不稳定、极端页码 OFFSET 溢出、超大请求体状态码错误及部分存储层错误被吞掉的问题。
## [0.3.2] - 2026-08-20
### Added
agent/blackboard_inheritance_test.go
+77 -1
View File
@@ -5,6 +5,7 @@ import (
"encoding/json"
"fmt"
"testing"
"time"
"github.com/Autumn-27/artex/db"
actool "github.com/Autumn-27/norma/tool"
@@ -23,6 +24,73 @@ func callReadJSON(t *testing.T, tool actool.CoreTool, input string) any {
return out
}
func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
d := testDB(t)
defer d.Close()
companies := d.Companies()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("overview-scope-%d", time.Now().UnixNano()), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = companies.DeleteCompany(companyID) })
domain := fmt.Sprintf("overview-scope-%d.invalid", companyID)
ip := fmt.Sprintf("2001:db8:%x::42", companyID%0xffff)
cidr := fmt.Sprintf("2001:db8:%x:1::/64", companyID%0xffff)
icp := fmt.Sprintf("京 ICP 备 %d 号", companyID)
keyword := fmt.Sprintf("Scope Company %d", companyID)
inputs := []db.ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "ip", Value: ip},
{Kind: "cidr", Value: cidr},
{Kind: "icp", Value: icp},
{Kind: "keyword", Value: keyword},
}
added, skipped, invalid, scopeErrors := companies.AddScopeInputs(companyID, inputs, "task context test")
if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 {
t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors)
}
task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
tools := NewToolSet(d.Exploration(task.ExplorationID), "planner")
tools.SetTaskID(task.ID)
tools.SetAssetStore(d.Assets(), companies)
overview := tools.graphOverviewData()
coverage, ok := overview["coverage"].(map[string]any)
if !ok {
t.Fatalf("coverage missing: %#v", overview["coverage"])
}
scopeRows, ok := coverage["scope"].([]map[string]any)
if !ok || len(scopeRows) != 1 {
t.Fatalf("task scope missing: %#v", coverage["scope"])
}
companyScope, ok := scopeRows[0]["company_scope"].([]map[string]any)
if !ok || len(companyScope) != len(inputs) {
t.Fatalf("company scope not expanded: %#v", scopeRows[0])
}
kinds := make(map[string]string, len(companyScope))
for _, rule := range companyScope {
kinds[fmt.Sprint(rule["kind"])] = fmt.Sprint(rule["value"])
}
for _, input := range inputs {
if kinds[input.Kind] != input.Value {
t.Errorf("scope %s=%q want %q", input.Kind, kinds[input.Kind], input.Value)
}
}
keywords, ok := scopeRows[0]["company_keywords"].([]string)
if !ok || len(keywords) != 1 || keywords[0] != keyword {
t.Fatalf("company keywords missing: %#v", scopeRows[0]["company_keywords"])
}
}
func TestBlackboardToolsReadDirectSources(t *testing.T) {
d := testDB(t)
defer d.Close()
@@ -125,7 +193,15 @@ func TestBlackboardToolsReadDirectSources(t *testing.T) {
if related[0]["source_task_id"] != source.ID || related[0]["inherited"] != true {
t.Fatalf("related source provenance: %#v", related[0])
}
if facts, ok := related[0]["recent_facts"].([]map[string]any); !ok || len(facts) == 0 || facts[0]["id"] != sourceFact {
recentFacts, ok := related[0]["recent_facts"].([]map[string]any)
foundSourceFact := false
for _, fact := range recentFacts {
if fact["id"] == sourceFact {
foundSourceFact = true
break
}
}
if !ok || !foundSourceFact {
t.Fatalf("related fact summary: %#v", related[0]["recent_facts"])
}
if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding {
agent/capture.go
+16 -2
View File
@@ -66,6 +66,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
tbuf.WriteString(text)
}
lastTool := &runTrace{startedAt: time.Now()}
var lastUsage *llm.Usage
var finalText string
var rerr error
@@ -74,10 +75,10 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
flush()
if ctx.Err() != nil { // engine/user cancellation, not a provider failure
sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool)
rec(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail})
rec(activityWithUsage(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}, lastUsage))
return finalText, reason, ctx.Err()
}
rec(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()})
rec(activityWithUsage(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()}, lastUsage))
return finalText, reason, err
}
switch ev.Kind {
@@ -111,6 +112,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
// buffered final-answer text must stay for the KindResult de-dup.
if ev.Usage != nil {
u := *ev.Usage
lastUsage = &u
rec(db.Activity{Kind: "usage",
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
@@ -145,6 +147,18 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
return finalText, reason, rerr
}
func activityWithUsage(activity db.Activity, usage *llm.Usage) db.Activity {
if usage == nil {
return activity
}
u := *usage
activity.InputTokens = &u.InputTokens
activity.OutputTokens = &u.OutputTokens
activity.CacheReadTokens = &u.CacheReadTokens
activity.CacheWriteTokens = &u.CacheWriteTokens
return activity
}
// blocksText concatenates the text of a tool-result's content blocks.
func blocksText(blocks []llm.ContentBlock) string {
var b strings.Builder
agent/capture_usage_test.go
+87
View File
@@ -0,0 +1,87 @@
package agent
import (
"context"
"errors"
"iter"
"testing"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/agentcore"
"github.com/Autumn-27/norma/llm"
)
type captureUsageProvider struct {
stream func(context.Context, func(llm.StreamEvent, error) bool)
}
func (p captureUsageProvider) Stream(ctx context.Context, _ llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) { p.stream(ctx, yield) }
}
func TestCaptureRunPersistsUsageOnProviderFailure(t *testing.T) {
wantErr := errors.New("provider failed after reporting usage")
provider := captureUsageProvider{stream: func(_ context.Context, yield func(llm.StreamEvent, error) bool) {
yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 11, CacheReadTokens: 3}}, nil)
yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "partial"}, nil)
yield(llm.StreamEvent{Type: llm.SEMessageDelta, Usage: llm.Usage{OutputTokens: 7, CacheWriteTokens: 2}}, nil)
yield(llm.StreamEvent{}, wantErr)
}}
var activities []db.Activity
_, _, err := captureRun(context.Background(), agentcore.Options{Provider: provider, MaxTurns: 1}, "test", func(a db.Activity) {
activities = append(activities, a)
})
if !errors.Is(err, wantErr) {
t.Fatalf("captureRun error=%v, want %v", err, wantErr)
}
assertCapturedResultUsage(t, activities, 11, 7, 3, 2)
}
func TestCaptureRunPersistsUsageOnCancellation(t *testing.T) {
started := make(chan struct{})
provider := captureUsageProvider{stream: func(ctx context.Context, yield func(llm.StreamEvent, error) bool) {
yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 13, CacheReadTokens: 5}}, nil)
yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "partial"}, nil)
yield(llm.StreamEvent{Type: llm.SEMessageDelta, Usage: llm.Usage{OutputTokens: 9, CacheWriteTokens: 4}}, nil)
close(started)
<-ctx.Done()
yield(llm.StreamEvent{}, ctx.Err())
}}
ctx, cancel := context.WithCancel(context.Background())
var activities []db.Activity
done := make(chan error, 1)
go func() {
_, _, err := captureRun(ctx, agentcore.Options{Provider: provider, MaxTurns: 1}, "test", func(a db.Activity) {
activities = append(activities, a)
})
done <- err
}()
<-started
cancel()
if err := <-done; !errors.Is(err, context.Canceled) {
t.Fatalf("captureRun error=%v, want context canceled", err)
}
assertCapturedResultUsage(t, activities, 13, 9, 5, 4)
}
func assertCapturedResultUsage(t *testing.T, activities []db.Activity, input, output, read, write int) {
t.Helper()
results := 0
for _, activity := range activities {
if activity.Kind != "result" {
continue
}
results++
if activity.InputTokens == nil || *activity.InputTokens != input ||
activity.OutputTokens == nil || *activity.OutputTokens != output ||
activity.CacheReadTokens == nil || *activity.CacheReadTokens != read ||
activity.CacheWriteTokens == nil || *activity.CacheWriteTokens != write {
t.Fatalf("result usage=%+v, want input=%d output=%d read=%d write=%d", activity, input, output, read, write)
}
}
if results != 1 {
t.Fatalf("result activity count=%d, activities=%+v", results, activities)
}
}
agent/tools.go
+36
View File
@@ -392,6 +392,42 @@ func (t *ToolSet) graphOverviewData() map[string]any {
e["value"] = r.Net
case r.CompanyID != nil:
e["company_id"] = *r.CompanyID
if t.cs != nil {
if company, err := t.cs.GetCompany(*r.CompanyID); err == nil && company != nil {
e["company_name"] = company.Name
}
if rules, err := t.cs.GetScope(*r.CompanyID); err == nil {
keywords := make([]string, 0)
companyScope := make([]map[string]any, 0, len(rules))
for _, rule := range rules {
value := rule.Raw
if value == "" {
switch rule.Kind {
case "domain":
value = rule.Domain
case "ip", "cidr":
value = rule.Net
default:
value = rule.Value
}
}
entry := map[string]any{"kind": rule.Kind, "value": value}
if rule.Reason != "" {
entry["reason"] = rule.Reason
}
companyScope = append(companyScope, entry)
if rule.Kind == "keyword" && rule.Raw != "" {
keywords = append(keywords, rule.Raw)
}
}
if len(companyScope) > 0 {
e["company_scope"] = companyScope
}
if len(keywords) > 0 {
e["company_keywords"] = keywords
}
}
}
}
scope = append(scope, e)
}
db/asset_dsl.go
+16 -1
View File
@@ -457,6 +457,18 @@ func isIdentChar(c byte) bool {
// ── QueryDSL ─────────────────────────────────────────────────────────────────
// ValidateDSL parses and compiles a DSL expression without touching the
// database. HTTP callers use it to distinguish client syntax errors from query
// failures, which must remain server errors.
func ValidateDSL(dsl string) error {
node, err := ParseDSL(dsl)
if err != nil {
return err
}
_, _, err = buildDSLWhere(node)
return err
}
// CountDSL returns the total number of assets matching a DSL expression (and optional
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
@@ -483,6 +495,9 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
node, err := ParseDSL(dsl)
if err != nil {
return nil, err
@@ -497,7 +512,7 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
}
args = append(args, limit, offset)
q := assetSelectCols + " WHERE " + where +
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
fmt.Sprintf(" ORDER BY last_seen DESC, id DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
db/assets.go
+195 -45
View File
@@ -83,6 +83,7 @@ type PortService struct {
type AssetStore struct {
db *DB
company *CompanyStore
tx *sql.Tx
}
// Assets returns the asset store.
@@ -93,6 +94,39 @@ func (d *DB) Assets() *AssetStore {
// Companies returns the company store associated with this asset store.
func (s *AssetStore) Companies() *CompanyStore { return s.company }
// withCompanyScopeMutation serializes scope resolution and every asset write
// that consumes its result in one transaction. Nested asset side effects reuse
// the same transaction through the scoped store.
func (s *AssetStore) withCompanyScopeMutation(fn func(*AssetStore) (int64, error)) (int64, error) {
if s.tx != nil {
return fn(s)
}
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
id, err := fn(scoped)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return id, nil
}
func (s *AssetStore) resolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
if s.tx != nil {
return resolveCompanyWithICP(s.tx, rootDomain, ipStr, icp)
}
return s.company.ResolveCompanyWithICP(rootDomain, ipStr, icp)
}
// =====================================================================
// Helpers
// =====================================================================
@@ -219,7 +253,12 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
companyID, err := s.company.ResolveCompany(domain, "")
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertRootDomain(req)
})
}
companyID, err := s.resolveCompanyWithICP(domain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -237,12 +276,20 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, $4::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, company_source, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, 'scope', $4::bigint[])
ON CONFLICT (domain) WHERE type = 'root_domain' DO UPDATE SET
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
extra = assets.extra || EXCLUDED.extra,
last_seen = now()
@@ -267,8 +314,13 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
if req.IP == "" {
return 0, fmt.Errorf("ip is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertIP(req)
})
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany("", req.IP)
companyID, err := s.resolveCompanyWithICP("", req.IP, "")
if err != nil {
return 0, err
}
@@ -292,9 +344,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, company_source, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, 'scope', $6::bigint[])
ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
bound_domains = (SELECT ARRAY(SELECT DISTINCT unnest(assets.bound_domains || EXCLUDED.bound_domains))),
open_ports = (
@@ -306,7 +358,15 @@ ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
)
),
c_segment = COALESCE(assets.c_segment, EXCLUDED.c_segment),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.IP, csegVal, boundDomains, openPortsJSON, companyID, taskIDs).Scan(&id)
@@ -367,14 +427,19 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertSubdomain(req)
})
}
rootDomain, _ := RootDomain(domain)
if rootDomain == "" {
rootDomain = domain
}
// resolve company by root domain
companyID, err := s.company.ResolveCompany(rootDomain, "")
// Resolve by root domain first, then the asset's exact normalized ICP.
companyID, err := s.resolveCompanyWithICP(rootDomain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -409,7 +474,7 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
cseg := calcCSegment(ipStr)
if companyID == nil && ipStr != "" {
companyID, _ = s.company.ResolveCompany("", ipStr)
companyID, _ = s.resolveCompanyWithICP("", ipStr, "")
}
var taskIDs string
@@ -432,14 +497,22 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
recordType := req.RecordType
recordValueArr := marshalStringArray(req.RecordValue)
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, $9::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, company_source, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, 'scope', $9::bigint[])
ON CONFLICT (domain, COALESCE(record_type,'')) WHERE type = 'subdomain' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
record_value = (SELECT ARRAY(SELECT DISTINCT unnest(assets.record_value || EXCLUDED.record_value))),
extra = assets.extra || EXCLUDED.extra,
@@ -459,7 +532,7 @@ type UpsertAppReq struct {
Category string
Description string
ICP string
CompanyID *int64 // explicit override; nil = no auto attribution for apps
CompanyID *int64 // explicit override; nil = exact ICP auto-attribution when available
TaskID int64
}
@@ -468,6 +541,11 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
if req.Name == "" {
return 0, fmt.Errorf("app name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertApp(req)
})
}
var taskIDs string
if req.TaskID > 0 {
@@ -477,6 +555,7 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
var bundleVal, catVal, descVal, icpVal, companyIDVal any
companySource := "scope"
if req.BundleID != "" {
bundleVal = req.BundleID
}
@@ -491,36 +570,65 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
if req.CompanyID != nil {
companyIDVal = *req.CompanyID
companySource = "explicit"
} else {
companyID, err := s.resolveCompanyWithICP("", "", req.ICP)
if err != nil {
return 0, err
}
if companyID != nil {
companyIDVal = *companyID
}
}
var id int64
var err error
if req.BundleID != "" {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7, $8::bigint[])
ON CONFLICT (bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL DO UPDATE SET
app_name = COALESCE(EXCLUDED.app_name, assets.app_name),
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
} else {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6, $7::bigint[])
ON CONFLICT (app_name) WHERE type = 'app' AND bundle_id IS NULL DO UPDATE SET
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
}
return id, err
}
@@ -548,6 +656,11 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
if req.URL == "" {
return 0, fmt.Errorf("url is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertHTTPService(req)
})
}
normURL := normalizeURL(req.URL)
domain, port, serviceName := parseURL(normURL)
rootDomain, _ := RootDomain(domain)
@@ -556,7 +669,7 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -592,16 +705,16 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, service_type, service_name, domain, ip, port, root_domain,
c_segment, favicon_mmh3, technologies, status_code, content_length, page_title,
auth, company_id, task_ids
auth, company_id, company_source, task_ids
)
VALUES (
'service', $1, 'http', $2, $3, $4, $5, $6,
$7::cidr, $8, $9::text[], $10, $11, $12,
$13::jsonb[], $14, $15::bigint[]
$13::jsonb[], $14, 'scope', $15::bigint[]
)
ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
status_code = COALESCE(EXCLUDED.status_code, assets.status_code),
@@ -618,7 +731,15 @@ ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, serviceName, domainVal, ipVal, nullableInt(port), rootDomain,
@@ -675,6 +796,11 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
if req.ServiceName == "" {
return 0, fmt.Errorf("service_name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertOtherService(req)
})
}
// normalize service_name (lowercase) so the (domain,ip,port,service_name)
// dedup key doesn't split "SSH" and "ssh" into separate rows.
serviceName := strings.ToLower(strings.TrimSpace(req.ServiceName))
@@ -689,7 +815,7 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -721,12 +847,12 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, service_type, service_name, domain, ip, port, root_domain,
c_segment, auth, company_id, task_ids
c_segment, auth, company_id, company_source, task_ids
)
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, $9::bigint[])
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, 'scope', $9::bigint[])
ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other' DO UPDATE SET
domain = COALESCE(EXCLUDED.domain, assets.domain),
ip = COALESCE(EXCLUDED.ip, assets.ip),
@@ -738,7 +864,15 @@ ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE typ
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
serviceName, domainVal, ipVal, req.Port, rootDomainVal,
@@ -799,6 +933,11 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
if req.Method == "" {
return 0, fmt.Errorf("method is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertEndpoint(req)
})
}
method := strings.ToUpper(req.Method)
normURL := normalizeURL(req.URL)
domain, port, _ := parseURL(normURL)
@@ -808,7 +947,7 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -840,12 +979,12 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, method, domain, ip, port, root_domain,
params, company_id, task_ids, c_segment
params, company_id, company_source, task_ids, c_segment
)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, $9::bigint[], $10::cidr)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, 'scope', $9::bigint[], $10::cidr)
ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
params = (
SELECT ARRAY(
@@ -857,7 +996,15 @@ ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, method, domainVal, ipVal, nullableInt(port), rootDomainVal,
@@ -884,6 +1031,9 @@ func (s *AssetStore) QueryByType(typ string, limit, offset int) ([]*Asset, error
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
rows, err := s.db.Query(`
SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
@@ -897,7 +1047,7 @@ SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(method,''), array_to_json(params)::text, extra, last_seen::text
FROM assets
WHERE type = $1
ORDER BY last_seen DESC
ORDER BY last_seen DESC, id DESC
LIMIT $2 OFFSET $3`, typ, limit, offset)
if err != nil {
return nil, err
@@ -954,7 +1104,7 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) {
}
func pageClause(args *[]any, limit, offset int) string {
q := ` ORDER BY last_seen DESC`
q := ` ORDER BY last_seen DESC, id DESC`
if limit > 0 {
*args = append(*args, limit)
q += fmt.Sprintf(` LIMIT $%d`, len(*args))
db/assets_test.go
+50
View File
@@ -4,6 +4,7 @@ import (
"fmt"
"slices"
"testing"
"time"
)
// testSetup opens a DB and returns both stores. Skips if no PG.
@@ -896,3 +897,52 @@ func TestNormalizeURL(t *testing.T) {
}
}
}
func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
d, assets, _ := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
taskID := stamp
marker := fmt.Sprintf("stable-page-%d", stamp)
sharedSeen := time.Date(2026, time.August, 21, 9, 0, 0, 0, time.UTC)
ids := make([]int64, 0, 5)
for i := 0; i < 5; i++ {
var id int64
domain := fmt.Sprintf("%s-%d.invalid", marker, i)
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,task_ids,last_seen)
VALUES ('root_domain',$1,$1,ARRAY[$2]::bigint[],$3) RETURNING id`,
domain, taskID, sharedSeen).Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
defer func() { _, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[])`, ids) }()
want := slices.Clone(ids)
slices.Reverse(want)
collect := func(query func(limit, offset int) ([]*Asset, error)) []int64 {
t.Helper()
var got []int64
for offset := 0; offset < len(ids); offset += 2 {
page, err := query(2, offset)
if err != nil {
t.Fatal(err)
}
for _, asset := range page {
got = append(got, asset.ID)
}
}
return got
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryByTask(taskID, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("task pagination order=%v want=%v", got, want)
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryDSL(marker, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("DSL pagination order=%v want=%v", got, want)
}
}
db/companies.go
+457 -109
View File
@@ -2,9 +2,11 @@ package db
import (
"database/sql"
"errors"
"fmt"
"net"
"strings"
"unicode/utf8"
)
// =====================================================================
@@ -25,7 +27,7 @@ type Company struct {
type CompanyWithScope struct {
Company
Scope []ScopeRule `json:"scope"`
AssetCount int `json:"asset_count"`
AssetCount int `json:"asset_count"`
}
// ScopeRule is one company_scope row.
@@ -35,6 +37,7 @@ type ScopeRule struct {
Kind string `json:"kind"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Value string `json:"value,omitempty"`
Raw string `json:"raw"`
Reason string `json:"reason,omitempty"`
}
@@ -42,6 +45,51 @@ type ScopeRule struct {
// CompanyStore operates on the companies + company_scope tables.
type CompanyStore struct{ db *DB }
var (
ErrCompanyNameConflict = errors.New("company name already exists")
ErrCompanyNotFound = errors.New("company not found")
)
const (
// MaxCompanyScopeRules bounds both request processing and the total number
// of persisted rules for one company.
MaxCompanyScopeRules = 256
// Raw and normalized textual scope payloads are bounded by Unicode rune
// count so multi-byte input is treated consistently by the API and DB layer.
MaxCompanyScopeRawRunes = 1024
MaxCompanyScopeValueRunes = 1024
)
// CompanyScopeValidationError identifies a client-correctable scope error.
// Storage and transaction failures are returned as ordinary errors instead.
type CompanyScopeValidationError struct{ Message string }
func (e *CompanyScopeValidationError) Error() string { return e.Message }
// ValidateCompanyScopeInputBounds applies request-wide limits before parsing.
// Store methods call it again so non-HTTP callers cannot bypass the limits.
func ValidateCompanyScopeInputBounds(inputs []ScopeInput) error {
if len(inputs) > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 最多 %d 条", MaxCompanyScopeRules,
)}
}
for i, input := range inputs {
if utf8.RuneCountInString(input.Value) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
}
return nil
}
// Scope writes rebuild derived asset ownership globally, so serialize them to
// ensure the committed attribution always reflects the latest committed rules.
// This key is reserved for company mutations; 7337741001 is the schema lock and
// 7337741002 is the cross-package test-suite lock.
const companyScopeMutationLock int64 = 7337741003
// Companies returns the company store.
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
@@ -69,6 +117,60 @@ RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
return
}
// CreateCompanyWithScope creates a company without updating an existing row.
// The company, its valid initial scope rules, and derived asset attribution are
// committed atomically. Invalid inputs retain the legacy partial-validation
// contract and are reported without preventing valid rules from being stored.
func (s *CompanyStore) CreateCompanyWithScope(name, logo string, inputs []ScopeInput, reason string) (
id int64, added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, 0, nil, err
}
rules, invalid, validationErrors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
nkey := companyNKey(name)
var logoVal any
if logo != "" {
logoVal = logo
}
if err := tx.QueryRow(`
INSERT INTO companies(name, nkey, logo)
VALUES ($1, $2, $3)
ON CONFLICT (nkey) DO NOTHING
RETURNING id`, name, nkey, logoVal).Scan(&id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return 0, 0, 0, invalid, validationErrors, ErrCompanyNameConflict
}
return 0, 0, 0, invalid, validationErrors, err
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, id, rules, reason)
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
}
if err := tx.Commit(); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
return id, added, skipped, invalid, validationErrors, nil
}
// GetCompany returns one company by id (nil if not found).
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
c := &Company{}
@@ -102,12 +204,55 @@ func (s *CompanyStore) UpsertByName(name string) (int64, error) {
return id, err
}
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
// DeleteCompany deletes a company and re-evaluates automatic ownership against
// the remaining companies in the same transaction. Explicitly-owned assets are
// detached by the FK and may then fall back to a remaining scope match.
func (s *CompanyStore) DeleteCompany(id int64) error {
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
_, err := s.DeleteCompanyWithAssets(id, false)
return err
}
// DeleteCompanyWithAssets deletes a company and optionally all of its assets in
// one transaction, then re-evaluates ownership against the remaining companies.
func (s *CompanyStore) DeleteCompanyWithAssets(id int64, deleteAssets bool) (assetsDeleted int64, err error) {
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
if deleteAssets {
res, err := tx.Exec(`DELETE FROM assets WHERE company_id = $1`, id)
if err != nil {
return 0, err
}
assetsDeleted, err = res.RowsAffected()
if err != nil {
return 0, err
}
}
res, err := tx.Exec(`DELETE FROM companies WHERE id = $1`, id)
if err != nil {
return 0, err
}
deleted, err := res.RowsAffected()
if err != nil {
return 0, err
}
if deleted == 0 {
return 0, ErrCompanyNotFound
}
if err := recomputeAttributionTx(tx); err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return assetsDeleted, nil
}
// ListCompanies returns all companies with scope and asset count.
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
rows, err := s.db.Query(`
@@ -149,7 +294,7 @@ ORDER BY c.name`)
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
rows, err := s.db.Query(`
SELECT id, company_id, kind,
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
COALESCE(domain,''), COALESCE(net::text,''), COALESCE(value,''), raw, COALESCE(reason,'')
FROM company_scope
WHERE company_id = $1
ORDER BY id`, companyID)
@@ -157,10 +302,10 @@ ORDER BY id`, companyID)
return nil, err
}
defer rows.Close()
var out []ScopeRule
out := make([]ScopeRule, 0)
for rows.Next() {
var r ScopeRule
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Value, &r.Raw, &r.Reason); err != nil {
return nil, err
}
out = append(out, r)
@@ -171,48 +316,172 @@ ORDER BY id`, companyID)
// AddScope parses and inserts scope lines for a company, then reattributes assets.
// Returns counts of added, skipped, and invalid lines.
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
rule, err := ParseScopeLine(line)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
inserted, err := s.insertScopeRule(companyID, rule, reason)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
if inserted {
added++
} else {
skipped++
}
inputs = append(inputs, ScopeInput{Value: line})
}
// attribute newly-covered assets
if added > 0 {
_ = s.attributeByCompany(companyID)
}
return
return s.AddScopeInputs(companyID, inputs, reason)
}
// insertScopeRule inserts a scope rule. Returns (inserted, error):
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
// AddScopeInputs inserts structured scope rules. Empty kinds retain the legacy
// domain/IP/CIDR auto-detection used by AddScope.
func (s *CompanyStore) AddScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, skipped, invalid int, errors []string) {
added, skipped, invalid, validationErrors, err := s.AddScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, skipped, invalid, validationErrors
}
// AddScopeInputsChecked inserts structured scope rules while keeping input
// validation separate from storage and transaction errors.
func (s *CompanyStore) AddScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, nil, err
}
rules, invalid, errors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, invalid, errors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, invalid, errors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if len(rules) == 0 {
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return 0, 0, invalid, errors, nil
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, invalid, errors, fmt.Errorf("重新计算企业归属失败: %w", err)
}
}
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return added, skipped, invalid, errors, nil
}
func parseScopeInputs(inputs []ScopeInput) (rules []ParsedScope, invalid int, validationErrors []string) {
rules = make([]ParsedScope, 0, len(inputs))
for _, input := range inputs {
rule, err := ParseScopeInput(input)
if err != nil {
invalid++
validationErrors = append(validationErrors, fmt.Sprintf("%s: %v", input.Value, err))
continue
}
rules = append(rules, rule)
}
return rules, invalid, validationErrors
}
func validateParsedScopeBounds(rules []ParsedScope) error {
for i, rule := range rules {
if utf8.RuneCountInString(rule.Raw) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
if utf8.RuneCountInString(rule.Value) > MaxCompanyScopeValueRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条规范化值过长: 最多 %d 个字符", i+1, MaxCompanyScopeValueRunes,
)}
}
}
return nil
}
func ensureCompanyExistsTx(tx *sql.Tx, companyID int64) error {
var exists bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM companies WHERE id = $1)`, companyID).Scan(&exists); err != nil {
return err
}
if !exists {
return ErrCompanyNotFound
}
return nil
}
func ensureCompanyScopeCapacityTx(tx *sql.Tx, companyID int64) error {
var count int
if err := tx.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id = $1`, companyID).Scan(&count); err != nil {
return err
}
if count > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 每个企业最多 %d 条", MaxCompanyScopeRules,
)}
}
return nil
}
func lockCompanyScopeMutation(tx *sql.Tx) error {
_, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, companyScopeMutationLock)
return err
}
func insertScopeRulesTx(tx *sql.Tx, companyID int64, rules []ParsedScope, reason string) (
added, skipped int, needsAttribution bool, err error,
) {
for _, rule := range rules {
inserted, insertErr := insertScopeRuleTx(tx, companyID, rule, reason)
if insertErr != nil {
return 0, 0, false, insertErr
}
if !inserted {
skipped++
continue
}
added++
needsAttribution = needsAttribution || rule.Kind != "keyword"
}
return added, skipped, needsAttribution, nil
}
// insertScopeRuleTx inserts a scope rule. inserted=false means a duplicate was
// ignored by ON CONFLICT, not an error.
func insertScopeRuleTx(tx *sql.Tx, companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
var res interface{ RowsAffected() (int64, error) }
if rule.Kind == "domain" {
res, err = s.db.Exec(`
switch rule.Kind {
case "domain":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, domain, raw, reason)
VALUES ($1, 'domain', $2, $3, $4)
ON CONFLICT ON CONSTRAINT uq_sv2_domain DO NOTHING`,
companyID, rule.Domain, rule.Raw, reason)
} else {
res, err = s.db.Exec(`
case "ip", "cidr":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, net, raw, reason)
VALUES ($1, $2, $3::cidr, $4, $5)
ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
companyID, rule.Kind, rule.Net, rule.Raw, reason)
case "icp", "keyword":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, value, raw, reason)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (company_id, kind, value) WHERE kind IN ('icp','keyword') DO NOTHING`,
companyID, rule.Kind, rule.Value, rule.Raw, reason)
default:
return false, fmt.Errorf("unsupported company scope kind %q", rule.Kind)
}
if err != nil {
return false, err
@@ -221,114 +490,179 @@ ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
return n > 0, nil
}
// RecomputeAttribution clears all company attribution on assets and rebuilds
// from scope rules. domain rules use root_domain equality; ip/cidr use INET >>= .
// RecomputeAttribution rebuilds only scope-derived ownership. Explicit company
// links are immutable under scope edits. Precedence is domain, IP/CIDR, then
// normalized exact ICP; keyword rules never attribute assets.
func (s *CompanyStore) RecomputeAttribution() error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// step 1: clear all attribution
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL`); err != nil {
if err := lockCompanyScopeMutation(tx); err != nil {
return err
}
// step 2: domain-based attribution (root_domain exact match)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind = 'domain'
AND a.root_domain = cs.domain
ORDER BY length(cs.domain) DESC
LIMIT 1
)
WHERE type IN ('root_domain','subdomain','service','endpoint')
AND root_domain IS NOT NULL`); err != nil {
if err := recomputeAttributionTx(tx); err != nil {
return err
}
// step 3: ip/cidr attribution (only for assets not yet attributed)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
ORDER BY masklen(cs.net) DESC
LIMIT 1
)
WHERE type IN ('ip','subdomain','service','endpoint')
AND ip IS NOT NULL
AND company_id IS NULL`); err != nil {
return err
}
return tx.Commit()
}
// attributeByCompany reattributes only assets matching a specific company's scope.
// Cheaper than full recompute; used when a company's scope is updated.
func (s *CompanyStore) attributeByCompany(companyID int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// clear this company's existing claims
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL WHERE company_id = $1`, companyID); err != nil {
return err
}
// domain rules
func recomputeAttributionTx(tx *sql.Tx) error {
// Only derived rows are cleared. Historical rows migrated without provenance
// are marked explicit by schema.sql, which is the non-destructive default.
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind = 'domain'
AND a.root_domain = cs.domain
AND type IN ('root_domain','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
UPDATE assets
SET company_id = NULL, company_source = 'scope'
WHERE company_source = 'scope'`); err != nil {
return err
}
// ip/cidr rules
// Domain-based attribution (root_domain exact match).
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
AND type IN ('ip','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'domain' AND a.root_domain = cs.domain
WHERE a.company_id IS NULL
AND a.type IN ('root_domain','subdomain','service','endpoint')
AND a.root_domain IS NOT NULL
ORDER BY a.id, length(cs.domain) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return tx.Commit()
// IP/CIDR attribution for still-unowned assets.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind IN ('ip','cidr') AND cs.net >>= a.ip::inet
WHERE a.company_id IS NULL
AND a.type IN ('ip','subdomain','service','endpoint')
AND a.ip IS NOT NULL
ORDER BY a.id, masklen(cs.net) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
// Exact normalized ICP attribution after domain/network precedence.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'icp'
AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = cs.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = cs.value
)
WHERE a.company_id IS NULL
AND (COALESCE(a.icp,'') <> '' OR COALESCE(a.app_icp,'') <> '')
ORDER BY a.id, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return nil
}
// UpdateScope replaces all scope rules for a company and reattributes.
func (s *CompanyStore) UpdateScope(companyID int64, lines []string, reason string) (added, invalid int, errs []string) {
// delete existing rules
if _, err := s.db.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
errs = append(errs, err.Error())
return
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
inputs = append(inputs, ScopeInput{Value: line})
}
added, _, invalid, errs = s.AddScope(companyID, lines, reason)
return
return s.UpdateScopeInputs(companyID, inputs, reason)
}
// UpdateScopeInputs replaces all rules with a structured set.
func (s *CompanyStore) UpdateScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, invalid int, errs []string) {
added, invalid, validationErrors, err := s.UpdateScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, invalid, validationErrors
}
// UpdateScopeInputsChecked replaces all rules while separating validation
// feedback from storage and transaction failures.
func (s *CompanyStore) UpdateScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, nil, err
}
rules, invalid, errs := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, invalid, errs, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, invalid, errs, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
if _, err := tx.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
return 0, invalid, errs, err
}
added, _, _, err = insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
// Rebuild even for an empty replacement because removing the old rules may
// detach scope-derived assets or expose a lower-precedence company match.
if err := recomputeAttributionTx(tx); err != nil {
return 0, invalid, errs, fmt.Errorf("重新计算企业归属失败: %w", err)
}
if err := tx.Commit(); err != nil {
return 0, invalid, errs, err
}
return added, invalid, errs, nil
}
// ResolveCompany returns the company_id for a given root_domain and/or ip, or nil
// if no scope rule matches. Mirrors the attribution logic used at asset insert time.
func (s *CompanyStore) ResolveCompany(rootDomain, ipStr string) (*int64, error) {
return s.ResolveCompanyWithICP(rootDomain, ipStr, "")
}
// ResolveCompanyWithICP mirrors RecomputeAttribution for insert-time ownership.
// ICP is consulted only after domain and IP/CIDR fail to match.
func (s *CompanyStore) ResolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
return resolveCompanyWithICP(s.db, rootDomain, ipStr, icp)
}
type companyScopeQueryer interface {
QueryRow(query string, args ...any) *sql.Row
}
func resolveCompanyWithICP(q companyScopeQueryer, rootDomain, ipStr, icp string) (*int64, error) {
if rootDomain != "" {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'domain'
AND domain = $1
ORDER BY length(domain) DESC
ORDER BY length(domain) DESC, company_id
LIMIT 1`, rootDomain).Scan(&cid)
if err == nil {
return &cid, nil
@@ -340,11 +674,11 @@ LIMIT 1`, rootDomain).Scan(&cid)
if ipStr != "" {
if net.ParseIP(ipStr) != nil {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind IN ('ip','cidr')
AND net >>= $1::inet
ORDER BY masklen(net) DESC
ORDER BY masklen(net) DESC, company_id
LIMIT 1`, ipStr).Scan(&cid)
if err == nil {
return &cid, nil
@@ -354,5 +688,19 @@ LIMIT 1`, ipStr).Scan(&cid)
}
}
}
if normalized := NormalizeICP(icp); normalized != "" {
var cid int64
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'icp' AND value = $1
ORDER BY company_id
LIMIT 1`, normalized).Scan(&cid)
if err == nil {
return &cid, nil
}
if err != sql.ErrNoRows {
return nil, err
}
}
return nil, nil
}
db/companies_test.go
+150
View File
@@ -1,7 +1,11 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
// cleanup helpers to remove test data
@@ -57,6 +61,70 @@ func TestCompanyUpsertAndGet(t *testing.T) {
}
}
func TestCreateCompanyWithScopeRejectsNormalizedDuplicate(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("Strict Company %d", stamp)
id, added, _, _, validationErrors, err := cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("strict-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
if added != 1 || len(validationErrors) != 0 {
t.Fatalf("initial scope: added=%d errors=%v", added, validationErrors)
}
_, _, _, _, _, err = cs.CreateCompanyWithScope(
" "+strings.ToUpper(strings.ReplaceAll(name, " ", " "))+" ",
"",
[]ScopeInput{{Kind: "domain", Value: fmt.Sprintf("replacement-%d.example", stamp)}},
"test",
)
if !errors.Is(err, ErrCompanyNameConflict) {
t.Fatalf("duplicate create error=%v want ErrCompanyNameConflict", err)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("strict-%d.example", stamp) {
t.Fatalf("duplicate create changed existing scope: %+v", scope)
}
}
func TestCreateCompanyWithScopeRollsBackOnScopeWriteFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Create %d", time.Now().UnixNano())
_, _, _, _, _, err = cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "test")
if err == nil {
t.Fatal("expected scope database write to fail")
}
company, getErr := cs.GetCompanyByName(name)
if getErr != nil {
t.Fatal(getErr)
}
if company != nil {
defer cleanupCompany(d, company.ID)
t.Fatalf("company row survived failed initial scope transaction: %+v", company)
}
}
func TestCompanyScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -189,6 +257,41 @@ func TestUpdateScope(t *testing.T) {
}
}
func TestUpdateScopeRollsBackOnInsertFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Scope Update %d", time.Now().UnixNano())
id, _, err := cs.UpsertCompany(name, "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
oldDomain := fmt.Sprintf("old-%d.example", time.Now().UnixNano())
added, _, invalid, addErrors := cs.AddScopeInputs(id, []ScopeInput{{Kind: "domain", Value: oldDomain}}, "initial")
if added != 1 || invalid != 0 || len(addErrors) != 0 {
t.Fatalf("seed scope: added=%d invalid=%d errors=%v", added, invalid, addErrors)
}
added, invalid, updateErrors := cs.UpdateScopeInputs(id, []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "replacement")
if added != 0 || invalid != 0 || len(updateErrors) == 0 {
t.Fatalf("failed update result: added=%d invalid=%d errors=%v", added, invalid, updateErrors)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != oldDomain {
t.Fatalf("failed replacement did not preserve old scope: %+v", scope)
}
}
func TestDeleteCompany(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -217,6 +320,53 @@ func TestDeleteCompany(t *testing.T) {
}
}
func TestDeleteCompanyWithAssetsDeletesBoth(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
id, _, err := cs.UpsertCompany(fmt.Sprintf("Delete Assets Company %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
var assetID int64
domain := fmt.Sprintf("delete-assets-%d.example", stamp)
if err := d.QueryRow(`
INSERT INTO assets(type, domain, root_domain, company_id, company_source)
VALUES ('root_domain', $1, $1, $2, 'explicit')
RETURNING id`, domain, id).Scan(&assetID); err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id = $1`, assetID) //nolint:errcheck
deleted, err := cs.DeleteCompanyWithAssets(id, true)
if err != nil {
t.Fatal(err)
}
if deleted != 1 {
t.Fatalf("assets deleted=%d want 1", deleted)
}
company, err := cs.GetCompany(id)
if err != nil {
t.Fatal(err)
}
if company != nil {
t.Fatalf("company still exists: %+v", company)
}
var assetsRemaining int
if err := d.QueryRow(`SELECT COUNT(*) FROM assets WHERE id = $1`, assetID).Scan(&assetsRemaining); err != nil {
t.Fatal(err)
}
if assetsRemaining != 0 {
t.Fatalf("asset %d survived company deletion", assetID)
}
}
func TestRecomputeAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
db/company_lock_test.go
+15
View File
@@ -0,0 +1,15 @@
package db
import "testing"
func TestCompanyScopeMutationLockDoesNotReuseInfrastructureLocks(t *testing.T) {
reserved := map[string]int64{
"schema migration": 7337741001,
"cross-package test suite": 7337741002,
}
for name, key := range reserved {
if companyScopeMutationLock == key {
t.Fatalf("company scope mutation lock reuses the %s advisory lock key %d", name, key)
}
}
}
db/company_scope.go
+61 -2
View File
@@ -4,19 +4,78 @@ import (
"fmt"
"net"
"strings"
"unicode"
"golang.org/x/net/publicsuffix"
)
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
// ParsedScope is one parsed asset-scope entry. Its kind selects Domain, Net, or Value.
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
type ParsedScope struct {
Kind string // "domain" | "ip" | "cidr"
Kind string // "domain" | "ip" | "cidr" | "icp" | "keyword"
Domain string // normalized registrable/root domain (kind=domain)
Net string // normalized CIDR, single IP as /32 or /128 (kind=ip|cidr)
Value string // normalized text (kind=icp|keyword)
Raw string // original input line
}
// ScopeInput is the structured API form for a company scope rule. Empty Kind
// preserves the legacy one-line auto detection for domain/IP/CIDR clients.
type ScopeInput struct {
Kind string `json:"kind,omitempty"`
Value string `json:"value"`
}
// NormalizeICP removes every Unicode whitespace character and folds case. ICP
// matching intentionally performs no fuzzy or punctuation normalization.
func NormalizeICP(value string) string {
return strings.ToLower(strings.Map(func(r rune) rune {
if unicode.IsSpace(r) {
return -1
}
return r
}, strings.TrimSpace(value)))
}
func normalizeKeyword(value string) string {
return strings.ToLower(strings.Join(strings.Fields(value), " "))
}
// ParseScopeInput validates an explicitly typed rule. Legacy callers can omit
// Kind and retain the existing domain/IP/CIDR auto-detection behavior.
func ParseScopeInput(input ScopeInput) (ParsedScope, error) {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
raw := strings.TrimSpace(input.Value)
if kind == "" {
return ParseScopeLine(raw)
}
switch kind {
case "domain", "ip", "cidr":
rule, err := ParseScopeLine(raw)
if err != nil {
return rule, err
}
if rule.Kind != kind {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("%q 不是有效的 %s 范围", raw, kind)
}
return rule, nil
case "icp":
value := NormalizeICP(raw)
if value == "" {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("ICP 不能为空")
}
return ParsedScope{Kind: kind, Value: value, Raw: raw}, nil
case "keyword":
value := normalizeKeyword(raw)
if value == "" {
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("企业关键词不能为空")
}
return ParsedScope{Kind: kind, Value: value, Raw: raw}, nil
default:
return ParsedScope{Kind: kind, Raw: raw}, fmt.Errorf("不支持的范围类型: %s", kind)
}
}
// ParseScopeLine classifies and validates one scope line (root domain / IP /
// CIDR). Guardrails reject bare TLDs and over-broad networks so a rule can never
// swallow the internet. IP ranges must be expressed as CIDR.
db/company_scope_consistency_test.go
+248
View File
@@ -0,0 +1,248 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
domain := fmt.Sprintf("scope-lock-%d.example", stamp)
oldCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock Old %d", stamp), "")
if err != nil {
t.Fatal(err)
}
newCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock New %d", stamp), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain = $1`, domain)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
})
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(oldCompany, []ScopeInput{
{Kind: "domain", Value: domain},
}, "test"); err != nil || added != 1 || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("seed old scope: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
mutation, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer mutation.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(mutation); err != nil {
t.Fatal(err)
}
if _, err := mutation.Exec(`DELETE FROM company_scope WHERE company_id=$1 AND kind='domain' AND domain=$2`, oldCompany, domain); err != nil {
t.Fatal(err)
}
if _, err := insertScopeRuleTx(mutation, newCompany, ParsedScope{
Kind: "domain", Domain: domain, Raw: domain,
}, "test"); err != nil {
t.Fatal(err)
}
started := make(chan struct{})
type result struct {
id int64
err error
}
resultCh := make(chan result, 1)
go func() {
close(started)
id, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
resultCh <- result{id: id, err: err}
}()
<-started
select {
case got := <-resultCh:
t.Fatalf("asset upsert escaped scope mutation lock: %+v", got)
case <-time.After(100 * time.Millisecond):
}
if err := mutation.Commit(); err != nil {
t.Fatal(err)
}
var got result
select {
case got = <-resultCh:
case <-time.After(3 * time.Second):
t.Fatal("asset upsert did not resume after scope mutation committed")
}
if got.err != nil {
t.Fatal(got.err)
}
var companyID *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, got.id).Scan(&companyID); err != nil {
t.Fatal(err)
}
if companyID == nil || *companyID != newCompany {
t.Fatalf("asset retained stale company: got=%v want=%d", companyID, newCompany)
}
}
func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
first, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break First %d", stamp), "")
if err != nil {
t.Fatal(err)
}
second, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break Second %d", stamp), "")
if err != nil {
t.Fatal(err)
}
want := min(first, second)
domain := fmt.Sprintf("tie-break-%d.example", stamp)
segment := uint64(stamp) & 0xffff
network := fmt.Sprintf("2001:db8:%x:1::/64", segment)
ip := fmt.Sprintf("2001:db8:%x:1::42", segment)
icp := fmt.Sprintf("ICP-TIE-%d", stamp)
appName := fmt.Sprintf("Tie Break App %d", stamp)
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
})
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "cidr", Value: network},
{Kind: "icp", Value: icp},
}
for _, companyID := range []int64{second, first} {
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != len(rules) || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("add company %d rules: added=%d invalid=%d validation=%v err=%v", companyID, added, invalid, validationErrors, err)
}
}
assertResolved := func(label string, got *int64, err error) {
t.Helper()
if err != nil {
t.Fatalf("%s resolve: %v", label, err)
}
if got == nil || *got != want {
t.Fatalf("%s resolve=%v want lowest company id %d", label, got, want)
}
}
got, err := companies.ResolveCompany(domain, "")
assertResolved("domain", got, err)
got, err = companies.ResolveCompany("", ip)
assertResolved("network", got, err)
got, err = companies.ResolveCompanyWithICP("", "", icp)
assertResolved("icp", got, err)
rootID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
if err != nil {
t.Fatal(err)
}
ipID, err := assets.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
appID, err := assets.UpsertApp(UpsertAppReq{Name: appName, ICP: icp})
if err != nil {
t.Fatal(err)
}
assertAssets := func(stage string) {
t.Helper()
for _, id := range []int64{rootID, ipID, appID} {
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, id).Scan(&got); err != nil {
t.Fatal(err)
}
if got == nil || *got != want {
t.Fatalf("%s asset %d company=%v want=%d", stage, id, got, want)
}
}
}
assertAssets("live")
if err := companies.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertAssets("recomputed")
}
func TestCompanyScopeLimitsAndCheckedErrors(t *testing.T) {
boundary := strings.Repeat("界", MaxCompanyScopeRawRunes)
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary}}); err != nil {
t.Fatalf("exact raw rune boundary rejected: %v", err)
}
var validationErr *CompanyScopeValidationError
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary + "界"}}); !errors.As(err, &validationErr) {
t.Fatalf("oversized raw value error=%v want CompanyScopeValidationError", err)
}
tooMany := make([]ScopeInput, MaxCompanyScopeRules+1)
if err := ValidateCompanyScopeInputBounds(tooMany); !errors.As(err, &validationErr) {
t.Fatalf("too many rules error=%v want CompanyScopeValidationError", err)
}
d, _, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Limits %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
rules := make([]ScopeInput, MaxCompanyScopeRules)
for i := range rules {
rules[i] = ScopeInput{Kind: "keyword", Value: fmt.Sprintf("limit-%d-%d", stamp, i)}
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != MaxCompanyScopeRules {
t.Fatalf("fill scope capacity: added=%d err=%v", added, err)
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, []ScopeInput{
{Kind: "keyword", Value: fmt.Sprintf("limit-%d-overflow", stamp)},
}, "test"); added != 0 || !errors.As(err, &validationErr) {
t.Fatalf("scope overflow: added=%d err=%v", added, err)
}
var count int
if err := d.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id=$1`, companyID).Scan(&count); err != nil {
t.Fatal(err)
}
if count != MaxCompanyScopeRules {
t.Fatalf("failed capacity mutation changed scope count=%d want=%d", count, MaxCompanyScopeRules)
}
missingID := companyID + 1_000_000_000
if _, _, _, _, err := companies.AddScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing add error=%v want ErrCompanyNotFound", err)
}
if _, _, _, err := companies.UpdateScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing update error=%v want ErrCompanyNotFound", err)
}
if _, err := companies.DeleteCompanyWithAssets(missingID, true); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing delete error=%v want ErrCompanyNotFound", err)
}
}
func TestCompanyScopeCheckedReturnsSystemErrorSeparately(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
companies := d.Companies()
if err := d.Close(); err != nil {
t.Fatal(err)
}
_, _, _, validationErrors, err := companies.AddScopeInputsChecked(1, nil, "test")
if err == nil {
t.Fatal("closed database did not return a system error")
}
if len(validationErrors) != 0 {
t.Fatalf("system error leaked into validation errors: %v", validationErrors)
}
var validationErr *CompanyScopeValidationError
if errors.As(err, &validationErr) || errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("system error misclassified: %v", err)
}
}
db/company_scope_test.go
+195 -2
View File
@@ -3,6 +3,7 @@ package db
import (
"fmt"
"testing"
"time"
)
// TestParseScopeLine covers classification + guardrails without a DB.
@@ -15,9 +16,9 @@ func TestParseScopeLine(t *testing.T) {
{"example.com", "domain", false},
{"https://sub.example.com/path", "domain", false},
{"1.2.3.4", "ip", false},
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"198.51.100.0/24", "cidr", false},
{"co.uk", "", true}, // bare public suffix
{"co.uk", "", true}, // bare public suffix
{"not a host", "", true},
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
}
@@ -39,6 +40,198 @@ func TestParseScopeLine(t *testing.T) {
}
}
func TestExplicitCompanyAttributionSurvivesScopeRebuild(t *testing.T) {
d, as, cs := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
explicitCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Explicit Attribution %d", stamp), "")
if err != nil {
t.Fatal(err)
}
autoCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Automatic Attribution %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, explicitCompany)
defer cleanupCompany(d, autoCompany)
domain := fmt.Sprintf("explicit-%d.invalid", stamp)
icp := fmt.Sprintf("ICP-%d", stamp)
network := fmt.Sprintf("2001:db8:%x::/64", uint64(stamp)&0xffff)
ip := fmt.Sprintf("2001:db8:%x::10", uint64(stamp)&0xffff)
// A pre-existing row with company_id and no provenance value represents old
// installations. The schema default conservatively treats it as explicit.
var assetID int64
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,company_id)
VALUES ('root_domain',$1,$1,$2) RETURNING id`, domain, explicitCompany).Scan(&assetID); err != nil {
t.Fatal(err)
}
defer deleteAsset(d, assetID)
appID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("explicit-app-%d", stamp), ICP: icp, CompanyID: &explicitCompany,
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, appID)
autoAssetID, err := as.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, autoAssetID)
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "icp", Value: icp},
{Kind: "cidr", Value: network},
}
if added, _, invalid, errs := cs.AddScopeInputs(autoCompany, rules, "test"); added != len(rules) || invalid != 0 {
t.Fatalf("AddScopeInputs: added=%d invalid=%d errors=%v", added, invalid, errs)
}
assertCompany := func(id, want int64, source string) {
t.Helper()
var got *int64
var gotSource string
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, id).Scan(&got, &gotSource); err != nil {
t.Fatal(err)
}
if got == nil || *got != want || gotSource != source {
t.Fatalf("asset %d company=%v source=%q, want %d/%q", id, got, gotSource, want, source)
}
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
assertCompany(autoAssetID, autoCompany, "scope")
// Replacing every matching rule detaches only the automatically-owned row.
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("replacement-%d.invalid", stamp)},
}, "test"); invalid != 0 || len(errs) != 0 {
t.Fatalf("UpdateScopeInputs: invalid=%d errors=%v", invalid, errs)
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
var autoCompanyID *int64
var autoSource string
if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, autoAssetID).Scan(&autoCompanyID, &autoSource); err != nil {
t.Fatal(err)
}
if autoCompanyID != nil || autoSource != "scope" {
t.Fatalf("automatic asset was not detached: company=%v source=%q", autoCompanyID, autoSource)
}
if err := cs.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertCompany(assetID, explicitCompany, "explicit")
assertCompany(appID, explicitCompany, "explicit")
// Deleting the explicitly selected company detaches through the FK, then the
// transactional rebuild may adopt the assets into a still-valid scope.
if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "icp", Value: icp},
}, "test"); invalid != 0 || len(errs) != 0 {
t.Fatalf("restore fallback scope: invalid=%d errors=%v", invalid, errs)
}
if err := cs.DeleteCompany(explicitCompany); err != nil {
t.Fatal(err)
}
assertCompany(assetID, autoCompany, "scope")
assertCompany(appID, autoCompany, "scope")
}
func TestParseStructuredCompanyScope(t *testing.T) {
icp, err := ParseScopeInput(ScopeInput{Kind: "ICP", Value: " 京ICP 备 123号-1\t"})
if err != nil {
t.Fatalf("parse ICP: %v", err)
}
if icp.Kind != "icp" || icp.Value != "京icp备123号-1" {
t.Fatalf("unexpected normalized ICP: %+v", icp)
}
keyword, err := ParseScopeInput(ScopeInput{Kind: "keyword", Value: " ACME Security "})
if err != nil {
t.Fatalf("parse keyword: %v", err)
}
if keyword.Value != "acme security" {
t.Fatalf("unexpected normalized keyword: %+v", keyword)
}
if _, err := ParseScopeInput(ScopeInput{Kind: "ip", Value: "example.com"}); err == nil {
t.Fatal("typed IP accepted a domain")
}
}
func TestCompanyICPAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
var suffix int64
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0)+1 FROM companies`).Scan(&suffix); err != nil {
t.Fatal(err)
}
cs := d.Companies()
as := d.Assets()
companyID, _, err := cs.UpsertCompany(fmt.Sprintf("ICP Scope Co %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE task_ids @> ARRAY[$1]::bigint[]`, suffix)
_, _ = d.Exec(`DELETE FROM companies WHERE id=$1`, companyID)
})
// A keyword can guide an Agent, but must never claim an asset by its name.
added, _, invalid, errs := cs.AddScopeInputs(companyID, []ScopeInput{
{Kind: "icp", Value: "京 ICP备 998877号"},
{Kind: "keyword", Value: "ICP Scope"},
}, "unit test")
if added != 2 || invalid != 0 || len(errs) != 0 {
t.Fatalf("add structured scope: added=%d invalid=%d errors=%v", added, invalid, errs)
}
rootID, err := as.UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("icp-scope-%d.example", suffix), ICP: "京icp备998877号", TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
appID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("ICP Scope Keyword Only %d", suffix), TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
icpAppID, err := as.UpsertApp(UpsertAppReq{
Name: fmt.Sprintf("ICP Matched App %d", suffix), ICP: " 京 ICP备 998877号 ", TaskID: suffix,
})
if err != nil {
t.Fatal(err)
}
assertCompany := func(assetID int64, want *int64) {
t.Helper()
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, assetID).Scan(&got); err != nil {
t.Fatal(err)
}
if want == nil && got != nil {
t.Fatalf("asset %d attributed by keyword: %d", assetID, *got)
}
if want != nil && (got == nil || *got != *want) {
t.Fatalf("asset %d company=%v want %d", assetID, got, *want)
}
}
assertCompany(rootID, &companyID)
assertCompany(appID, nil)
assertCompany(icpAppID, &companyID)
}
// TestCompanyScopeAttribution exercises the full loop against dev PG: create
// company (unique name), add scope, and verify auto-attribution at insert time,
// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
db/conversation.go
+47 -11
View File
@@ -50,19 +50,32 @@ GROUP BY c.id`)
// Conversation is one ChatGPT-style chat thread bound to an agent key. It lives
// independent of the pentest exploration graph — see schema.sql §I.
type Conversation struct {
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
Pinned bool `json:"pinned"`
PinnedAt *time.Time `json:"pinned_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at`
// ConversationPatch updates only the fields whose pointers are non-nil.
type ConversationPatch struct {
Title *string
Pinned *bool
}
const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at`
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
var c Conversation
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt)
var pinnedAt sql.NullTime
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt)
if pinnedAt.Valid {
c.Pinned = true
c.PinnedAt = &pinnedAt.Time
}
return c, err
}
@@ -127,7 +140,8 @@ func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error {
// ListConversations returns all threads, most-recently-updated first.
func (d *DB) ListConversations() ([]*Conversation, error) {
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`)
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations
ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
@@ -155,9 +169,31 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) {
return &c, nil
}
// RenameConversation sets a thread's title.
// UpdateConversation applies a partial title/pin mutation and returns the updated
// row. Pinning an already-pinned conversation preserves its original pin order.
func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`UPDATE conversations SET
title = CASE WHEN $2::boolean THEN $3 ELSE title END,
pinned_at = CASE
WHEN $4::boolean IS NULL THEN pinned_at
WHEN $4::boolean THEN COALESCE(pinned_at, now())
ELSE NULL
END
WHERE id=$1
RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &c, nil
}
// RenameConversation sets a thread's title. Kept for automatic first-message
// titles and compatibility with existing callers.
func (d *DB) RenameConversation(id int64, title string) error {
_, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title)
_, err := d.UpdateConversation(id, ConversationPatch{Title: &title})
return err
}
db/conversation_pin_test.go
+69
View File
@@ -0,0 +1,69 @@
package db
import (
"fmt"
"testing"
"time"
)
func TestConversationPinOrderingAndPatch(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
first, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-first-%d", suffix), nil)
if err != nil {
t.Fatal(err)
}
second, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-second-%d", suffix), nil)
if err != nil {
_ = d.DeleteConversation(first.ID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(first.ID)
_ = d.DeleteConversation(second.ID)
})
pinned := true
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || !first.Pinned || first.PinnedAt == nil {
t.Fatalf("pin first = %+v, %v", first, err)
}
firstPinnedAt := *first.PinnedAt
time.Sleep(5 * time.Millisecond)
second, err = d.UpdateConversation(second.ID, ConversationPatch{Pinned: &pinned})
if err != nil || second == nil || !second.Pinned {
t.Fatalf("pin second = %+v, %v", second, err)
}
newTitle := "renamed while pinned"
first, err = d.UpdateConversation(first.ID, ConversationPatch{Title: &newTitle, Pinned: &pinned})
if err != nil || first == nil {
t.Fatal(err)
}
if first.Title != newTitle || first.PinnedAt == nil || !first.PinnedAt.Equal(firstPinnedAt) {
t.Fatalf("repeat pin should preserve pin time: %+v (want %v)", first, firstPinnedAt)
}
conversations, err := d.ListConversations()
if err != nil {
t.Fatal(err)
}
positions := map[int64]int{}
for index, conversation := range conversations {
positions[conversation.ID] = index
}
if positions[second.ID] >= positions[first.ID] {
t.Fatalf("newer pin must sort first: second=%d first=%d", positions[second.ID], positions[first.ID])
}
pinned = false
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || first.Pinned || first.PinnedAt != nil {
t.Fatalf("unpin first = %+v, %v", first, err)
}
}
db/exploration.go
+215 -6
View File
@@ -3,11 +3,17 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"time"
)
// ErrIntentStateConflict marks a failed compare-and-set transition. Callers use
// it to distinguish an expected control race (HTTP 409) from a storage failure.
var ErrIntentStateConflict = errors.New("intent state changed concurrently")
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
@@ -70,6 +76,17 @@ type TokenUsage struct {
CacheWriteTokens int `json:"cache_write_tokens"`
}
// SessionTokenUsage is the authoritative completed-run total for one task UI
// session. Worker sessions are keyed by intent rather than the reusable work#N
// executor name, so retries and reassignment remain attached to the same row.
type SessionTokenUsage struct {
Session string `json:"session"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// DailyTokenBucket is one day's global token aggregate across all tasks.
type DailyTokenBucket struct {
Day string `json:"day"` // "YYYY-MM-DD"
@@ -264,6 +281,21 @@ WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, termin
return err
}
// CompareAndSetIntentState transitions one local intent only when its current
// state still matches expected. It is the state boundary used by pause/resume and
// worker settlement, so a stale API read cannot overwrite a concurrent finish.
func (s *ExplorationStore) CompareAndSetIntentState(id int64, expected, state string) (bool, error) {
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
res, err := s.db.Exec(`UPDATE exploration_nodes
SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $5 THEN now() ELSE NULL END
WHERE id=$2 AND exploration_id=$3 AND kind='intent' AND state=$4`, state, id, s.expID, expected, terminal)
if err != nil {
return false, err
}
n, err := res.RowsAffected()
return n == 1, err
}
// IntentCleanup summarizes the blackboard records removed when a user cancels
// one worker intent. Global assets and traffic are deliberately not part of this
// operation; exploration anchors disappear only because their owning nodes do.
@@ -295,7 +327,7 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, err
}
if state != "running" && state != "paused" {
return out, fmt.Errorf("intent state %s cannot be cancelled", state)
return out, fmt.Errorf("%w: intent state %s cannot be cancelled", ErrIntentStateConflict, state)
}
if err := tx.QueryRow(`SELECT
@@ -304,14 +336,33 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
FROM exploration_edges e
JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil {
AND n.kind IN ('fact','finding')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil {
return out, err
}
if _, err := tx.Exec(`DELETE FROM findings f USING exploration_edges e, exploration_nodes n
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.id=e.dst_id AND n.exploration_id=e.exploration_id
AND n.kind='finding' AND f.node_id=n.id`, s.expID, id, RelYields); err != nil {
AND n.kind='finding' AND f.node_id=n.id
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, s.expID, id, RelYields); err != nil {
return out, err
}
// Activity rows are part of the blackboard cleanup, but their token usage is
// irreversible metering data. Fold completed runs plus the latest unfinished
// usage frame into detached daily result rows before deleting the conversation.
// The rows are excluded from per-intent sessions while whole-task totals and
// the original UTC reporting dates remain accurate.
tokenBuckets, err := intentTokenRollup(tx, s.expID, id)
if err != nil {
return out, err
}
res, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id)
@@ -319,15 +370,36 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, err
}
out.Activities, _ = res.RowsAffected()
for _, bucket := range tokenBuckets {
metadata, _ := json.Marshal(map[string]any{
"cancelled_intent_id": id,
"token_day": bucket.Day.Format(time.DateOnly),
"token_rollup": true,
})
if _, err := tx.Exec(`INSERT INTO activity(
exploration_id, worker, kind, summary, metadata,
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at)
VALUES ($1,'token-ledger','result',$2,$3,$4,$5,$6,$7,$8)`,
s.expID, fmt.Sprintf("已取消意图 #%d 的 Token 计量", id), metadata,
bucket.Usage.InputTokens, bucket.Usage.OutputTokens,
bucket.Usage.CacheReadTokens, bucket.Usage.CacheWriteTokens, bucket.Day); err != nil {
return out, err
}
}
res, err = tx.Exec(`DELETE FROM exploration_nodes
WHERE exploration_id=$1 AND (
id=$2 OR id IN (
SELECT e.dst_id FROM exploration_edges e
JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')
)
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)
)
)`, s.expID, id, RelYields)
if err != nil {
return out, err
@@ -339,6 +411,108 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, tx.Commit()
}
type tokenUsageBucket struct {
Day time.Time
Usage TokenUsage
}
// intentTokenRollup returns exactly-once usage grouped by its original UTC day.
// Each result terminates one run. A result carrying usage is authoritative; an
// older error/cancel result without usage falls back to the latest cumulative
// usage frame since the previous result. A trailing frame represents a run that
// was interrupted before its terminal result could be persisted.
func intentTokenRollup(tx *sql.Tx, explorationID, intentID int64) ([]tokenUsageBucket, error) {
rows, err := tx.Query(`SELECT kind, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at
FROM activity
WHERE exploration_id=$1 AND node_id=$2 AND kind IN ('usage','result')
ORDER BY id`, explorationID, intentID)
if err != nil {
return nil, err
}
defer rows.Close()
byDay := make(map[time.Time]TokenUsage)
add := func(at time.Time, usage TokenUsage) {
if usage.InputTokens == 0 && usage.OutputTokens == 0 &&
usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
return
}
utc := at.UTC()
day := time.Date(utc.Year(), utc.Month(), utc.Day(), 0, 0, 0, 0, time.UTC)
total := byDay[day]
total.add(usage)
byDay[day] = total
}
var pending TokenUsage
var pendingAt time.Time
hasPending := false
for rows.Next() {
var kind string
var input, output, read, write sql.NullInt64
var createdAt time.Time
if err := rows.Scan(&kind, &input, &output, &read, &write, &createdAt); err != nil {
return nil, err
}
current := TokenUsage{
InputTokens: int(input.Int64),
OutputTokens: int(output.Int64),
CacheReadTokens: int(read.Int64),
CacheWriteTokens: int(write.Int64),
}
hasCurrent := input.Valid || output.Valid || read.Valid || write.Valid
if kind == "usage" {
if hasCurrent {
pending, pendingAt, hasPending = current, createdAt, true
}
continue
}
if hasCurrent {
// Preserve a partially populated legacy terminal row by filling only its
// missing dimensions from the latest cumulative usage frame.
if hasPending {
if !input.Valid {
current.InputTokens = pending.InputTokens
}
if !output.Valid {
current.OutputTokens = pending.OutputTokens
}
if !read.Valid {
current.CacheReadTokens = pending.CacheReadTokens
}
if !write.Valid {
current.CacheWriteTokens = pending.CacheWriteTokens
}
}
add(createdAt, current)
} else if hasPending {
add(createdAt, pending)
}
pending, pendingAt, hasPending = TokenUsage{}, time.Time{}, false
}
if err := rows.Err(); err != nil {
return nil, err
}
if hasPending {
add(pendingAt, pending)
}
buckets := make([]tokenUsageBucket, 0, len(byDay))
for day, usage := range byDay {
buckets = append(buckets, tokenUsageBucket{Day: day, Usage: usage})
}
sort.Slice(buckets, func(i, j int) bool { return buckets[i].Day.Before(buckets[j].Day) })
return buckets, nil
}
func (u *TokenUsage) add(other TokenUsage) {
u.InputTokens += other.InputTokens
u.OutputTokens += other.OutputTokens
u.CacheReadTokens += other.CacheReadTokens
u.CacheWriteTokens += other.CacheWriteTokens
}
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at`
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
@@ -761,6 +935,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
return out, rows.Err()
}
// TokenStatsBySession aggregates every persisted completed run, independent of
// activity-history pagination. Main/planner use fixed keys; Worker executions use
// their intent node id, even when a different work#N process handled a retry.
func (s *ExplorationStore) TokenStatsBySession() ([]SessionTokenUsage, error) {
rows, err := s.db.Query(`SELECT
CASE
WHEN COALESCE(a.worker,'')='mainagent' THEN 'main'
WHEN COALESCE(a.worker,'')='planner' THEN 'plan'
WHEN n.id IS NOT NULL THEN 'intent:' || n.id::text
ELSE ''
END AS session_key,
COALESCE(SUM(a.input_tokens),0), COALESCE(SUM(a.output_tokens),0),
COALESCE(SUM(a.cache_read_tokens),0), COALESCE(SUM(a.cache_write_tokens),0)
FROM activity a
LEFT JOIN exploration_nodes n
ON n.id=a.node_id AND n.exploration_id=a.exploration_id AND n.kind='intent'
WHERE a.exploration_id=$1 AND a.kind='result'
AND (COALESCE(a.worker,'') IN ('mainagent','planner') OR n.id IS NOT NULL)
GROUP BY session_key
ORDER BY session_key`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []SessionTokenUsage{}
for rows.Next() {
var u SessionTokenUsage
if err := rows.Scan(&u.Session, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out = append(out, u)
}
return out, rows.Err()
}
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
// Returns items and the new cursor (max id seen).
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
db/exploration_tokens_test.go
+130
View File
@@ -0,0 +1,130 @@
package db
import (
"fmt"
"testing"
)
func TestTokenStatsBySessionUsesCompleteIntentHistory(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("token sessions", "aggregate complete session history")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "worker session"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
factID, err := store.AddNode(KindFact, map[string]any{"summary": "not a worker"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
otherExpID, err := d.CreateExploration("other token sessions", "foreign intent must not match")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, otherExpID)
foreignIntentID, err := d.Exploration(otherExpID).AddIntent(map[string]any{"summary": "foreign worker"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
appendResult := func(nodeID *int64, worker string, input, output, read, write int) {
t.Helper()
if _, appendErr := store.AppendActivity(Activity{
NodeID: nodeID,
Worker: worker,
Kind: "result",
InputTokens: &input,
OutputTokens: &output,
CacheReadTokens: &read,
CacheWriteTokens: &write,
}); appendErr != nil {
t.Fatal(appendErr)
}
}
appendResult(nil, "mainagent", 11, 12, 13, 14)
appendResult(nil, "planner", 21, 22, 23, 24)
// More than ActivityPage's default limit proves this aggregate reads the
// persisted history directly instead of summing the currently loaded page.
const completedRuns = 205
for i := 0; i < completedRuns; i++ {
appendResult(&intentID, fmt.Sprintf("work#%d", i%3+1), 1, 2, 3, 4)
}
// All three rows remain part of the legacy worker/whole-task totals, but none
// represents a local Worker session and therefore none may create intent:*.
appendResult(&factID, "work#fact", 31, 32, 33, 34)
appendResult(nil, "work#missing-node", 41, 42, 43, 44)
appendResult(&foreignIntentID, "work#foreign", 51, 52, 53, 54)
ignoredInput := 1000
if _, err := store.AppendActivity(Activity{
NodeID: &intentID,
Worker: "work#1",
Kind: "usage",
InputTokens: &ignoredInput,
}); err != nil {
t.Fatal(err)
}
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
got := make(map[string]SessionTokenUsage, len(sessions))
for _, usage := range sessions {
got[usage.Session] = usage
}
if len(got) != 3 {
t.Fatalf("sessions = %+v, want only main, plan, and the local intent", sessions)
}
assertSessionTokens(t, got["main"], "main", 11, 12, 13, 14)
assertSessionTokens(t, got["plan"], "plan", 21, 22, 23, 24)
assertSessionTokens(t, got[fmt.Sprintf("intent:%d", intentID)], fmt.Sprintf("intent:%d", intentID),
completedRuns, completedRuns*2, completedRuns*3, completedRuns*4)
workers, err := store.TokenStatsByWorker()
if err != nil {
t.Fatal(err)
}
if len(workers) != 8 { // main, planner, 3 executors, and the 3 deliberately invalid session rows
t.Fatalf("legacy workers changed: got %d entries: %+v", len(workers), workers)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total,
11+21+completedRuns+31+41+51,
12+22+completedRuns*2+32+42+52,
13+23+completedRuns*3+33+43+53,
14+24+completedRuns*4+34+44+54)
}
func assertSessionTokens(t *testing.T, got SessionTokenUsage, session string, input, output, read, write int) {
t.Helper()
if got.Session != session || got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("session %q = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
session, got, input, output, read, write)
}
}
func assertTokenUsage(t *testing.T, got TokenUsage, input, output, read, write int) {
t.Helper()
if got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("token total = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
got, input, output, read, write)
}
}
db/findings.go
+200 -4
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
@@ -147,6 +148,11 @@ type FindingFilter struct {
Sort string // "severity" | "time"
}
// FindingUnassignedTask is the task filter sentinel for findings whose task is
// absent. That includes rows created without a task and rows retained after their
// originating task was deleted (the findings FK is ON DELETE SET NULL).
const FindingUnassignedTask = "__unassigned__"
// where builds the WHERE clause (shared by the page and count queries) plus its
// positional args. Only equality filters, all parameterized.
func (f FindingFilter) where() (string, []any) {
@@ -163,7 +169,9 @@ func (f FindingFilter) where() (string, []any) {
add("status", f.Status)
add("vulnclass", f.VulnClass)
// task_id 是 bigint 列,按整数比较(不能走上面的文本 add);空/非法值忽略。
if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
if f.TaskID == FindingUnassignedTask {
conds = append(conds, "(f.task_id IS NULL OR t.id IS NULL)")
} else if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
args = append(args, tid)
conds = append(conds, fmt.Sprintf("f.task_id = $%d", len(args)))
}
@@ -185,14 +193,19 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
where, args := f.where()
var total int
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f`+where, args...).Scan(&total); err != nil {
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f LEFT JOIN tasks t ON f.task_id=t.id`+where, args...).Scan(&total); err != nil {
return nil, 0, err
}
// Avoid overflowing (page-1)*pageSize for an arbitrarily large page number.
// Once the requested page is beyond the exact count, no data query is needed.
if total == 0 || page > (total-1)/pageSize+1 {
return []*DBFinding{}, total, nil
}
order := "f.created_at DESC"
order := "f.created_at DESC, f.id DESC"
if f.Sort == "severity" {
// critical > high > medium > low > 其它, then newest first.
order = `CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END DESC, f.created_at DESC`
order = `CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END DESC, f.created_at DESC, f.id DESC`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
q := fmt.Sprintf(`
@@ -210,6 +223,189 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
return out, total, err
}
// FindingGroup is one task-level bucket in the global findings view. TaskID is
// nil for both findings that never had a task and findings retained after task
// deletion; those records intentionally share one "unassigned/deleted" bucket.
type FindingGroup struct {
TaskID *int64 `json:"task_id"`
TaskDescription string `json:"task_description"`
TaskStatus string `json:"task_status"`
Count int `json:"count"`
Critical int `json:"critical"`
High int `json:"high"`
Medium int `json:"medium"`
Low int `json:"low"`
LastFoundAt time.Time `json:"last_found_at"`
}
// ListFindingGroups returns a page of task groups matching the same filters as
// ListFindingsPage. The group count and finding count are independent totals so
// clients can page groups without losing the exact export/selection count.
func (d *DB) ListFindingGroups(f FindingFilter, page, pageSize int) ([]FindingGroup, int, int, error) {
if page <= 0 {
page = 1
}
if pageSize <= 0 {
pageSize = 10
}
where, args := f.where()
grouped := ` FROM findings f LEFT JOIN tasks t ON f.task_id=t.id` + where +
` GROUP BY t.id, t.description, t.status, t.paused, t.queued`
var groupTotal, findingTotal int
countQuery := `SELECT COUNT(*), COALESCE(SUM(finding_count),0) FROM (` +
`SELECT COUNT(*) AS finding_count` + grouped + `) grouped_findings`
if err := d.QueryRow(countQuery, args...).Scan(&groupTotal, &findingTotal); err != nil {
return nil, 0, 0, err
}
if groupTotal == 0 || page > (groupTotal-1)/pageSize+1 {
return []FindingGroup{}, groupTotal, findingTotal, nil
}
order := "MAX(f.created_at) DESC, t.id DESC NULLS LAST"
if f.Sort == "severity" {
order = `MAX(CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END) DESC, MAX(f.created_at) DESC, t.id DESC NULLS LAST`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
query := fmt.Sprintf(`SELECT t.id, COALESCE(t.description,''), COALESCE(
CASE
WHEN t.status IN ('done','failed','timeout') THEN t.status
WHEN t.queued THEN 'queued'
WHEN t.paused THEN 'paused'
ELSE t.status
END, ''),
COUNT(*),
COUNT(*) FILTER (WHERE f.severity='critical'),
COUNT(*) FILTER (WHERE f.severity='high'),
COUNT(*) FILTER (WHERE f.severity='medium'),
COUNT(*) FILTER (WHERE f.severity='low'),
MAX(f.created_at)%s
ORDER BY %s LIMIT $%d OFFSET $%d`, grouped, order, len(args)+1, len(args)+2)
rows, err := d.Query(query, pageArgs...)
if err != nil {
return nil, 0, 0, err
}
defer rows.Close()
groups := []FindingGroup{}
for rows.Next() {
var group FindingGroup
var taskID sql.NullInt64
if err := rows.Scan(&taskID, &group.TaskDescription, &group.TaskStatus, &group.Count,
&group.Critical, &group.High, &group.Medium, &group.Low, &group.LastFoundAt); err != nil {
return nil, 0, 0, err
}
if taskID.Valid {
id := taskID.Int64
group.TaskID = &id
}
groups = append(groups, group)
}
return groups, groupTotal, findingTotal, rows.Err()
}
// ErrFindingOriginUnavailable means a retained finding no longer has a live
// owning task and finding node from which a follow-up intent can be derived.
var ErrFindingOriginUnavailable = errors.New("finding origin is no longer available")
// AddFindingFollowUpIntent atomically creates a priority-10 human intent from a
// live finding node, copies that finding's asset anchors, records the
// finding --derived_from--> intent lineage edge, and persists its audit activity.
// The returned activity is the committed row and can be broadcast as-is without
// calling AppendActivity again.
func (s *ExplorationStore) AddFindingFollowUpIntent(findingID, findingNodeID int64, description string, audit Activity) (int64, Activity, error) {
tx, err := s.db.Begin()
if err != nil {
return 0, Activity{}, err
}
defer tx.Rollback()
var liveNodeID int64
err = tx.QueryRow(`SELECT n.id
FROM findings f
JOIN tasks t ON t.id=f.task_id
JOIN exploration_nodes n ON n.id=f.node_id AND n.exploration_id=t.exploration_id
WHERE f.id=$1 AND f.node_id=$2 AND t.exploration_id=$3 AND n.kind='finding'
FOR SHARE OF f, t, n`, findingID, findingNodeID, s.expID).Scan(&liveNodeID)
if err == sql.ErrNoRows {
return 0, Activity{}, ErrFindingOriginUnavailable
}
if err != nil {
return 0, Activity{}, err
}
anchors := []int64{}
anchorRows, err := tx.Query(`SELECT asset_id FROM exploration_anchors WHERE node_id=$1 ORDER BY asset_id`, liveNodeID)
if err != nil {
return 0, Activity{}, err
}
for anchorRows.Next() {
var assetID int64
if err := anchorRows.Scan(&assetID); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
anchors = append(anchors, assetID)
}
if err := anchorRows.Err(); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
if err := anchorRows.Close(); err != nil {
return 0, Activity{}, err
}
payload := map[string]any{
"summary": description,
"source_finding_id": findingID,
"source_finding_node_id": liveNodeID,
}
if len(anchors) > 0 {
payload["asset_ids"] = anchors
}
raw, err := json.Marshal(payload)
if err != nil {
return 0, Activity{}, err
}
var intentID int64
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
VALUES ($1,'intent',$2,10,'open','human') RETURNING id`, s.expID, raw).Scan(&intentID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id)
SELECT $1, asset_id FROM exploration_anchors WHERE node_id=$2
ON CONFLICT DO NOTHING`, intentID, liveNodeID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id)
VALUES ($1,$2,$3,$4)`, s.expID, liveNodeID, RelDerivedFrom, intentID); err != nil {
return 0, Activity{}, err
}
audit.NodeID = &intentID
if summary := strings.TrimSpace(audit.Summary); summary != "" {
audit.Summary = fmt.Sprintf("%s #%d", summary, intentID)
} else {
audit.Summary = ""
}
metadata := audit.Metadata
if len(metadata) == 0 {
metadata = json.RawMessage(`{}`)
}
if err := tx.QueryRow(`
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14)
RETURNING id, created_at`, s.expID, audit.NodeID, utf8Clean(audit.Worker), utf8Clean(audit.Kind), utf8Clean(audit.Tool), utf8Clean(audit.ToolUseID), audit.IsError,
utf8Clean(audit.Summary), utf8Clean(audit.Detail), metadata, audit.InputTokens, audit.OutputTokens, audit.CacheReadTokens, audit.CacheWriteTokens).
Scan(&audit.ID, &audit.CreatedAt); err != nil {
return 0, Activity{}, err
}
audit.Metadata = metadata
if err := tx.Commit(); err != nil {
return 0, Activity{}, err
}
return intentID, audit, nil
}
// ListFindingsForExport returns findings for the 发现 page 导出功能，携带完整
// report 字段、不分页。ids 非空时按这批 finding id 精确导出(勾选导出),忽略
// filter;ids 为空时按 filter 导出(导出当前筛选/全部)。结果按严重等级降序、
db/findings_test.go
+298
View File
@@ -1,8 +1,12 @@
package db
import (
"encoding/json"
"errors"
"fmt"
"slices"
"testing"
"time"
)
// TestDeleteFinding verifies删除漏洞 removes both the findings row and its
@@ -183,3 +187,297 @@ func TestFindingsPageAndStats(t *testing.T) {
t.Fatalf("SetFindingSeverity(-1): want 0,nil got %d,%v", n, err)
}
}
func TestListFindingsPageUsesStableIDTieBreaker(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_stable_page_%d__", time.Now().UnixNano())
defer d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc)
ids := make([]int64, 0, 7)
for i := 0; i < 7; i++ {
id, addErr := d.AddFinding(0, 0, vc, "", SeverityHigh, fmt.Sprintf("finding %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
ids = append(ids, id)
}
sharedCreatedAt := time.Date(2026, time.August, 21, 8, 30, 0, 0, time.UTC)
if _, err := d.Exec(`UPDATE findings SET created_at=$1 WHERE vulnclass=$2`, sharedCreatedAt, vc); err != nil {
t.Fatal(err)
}
want := make([]int64, len(ids))
for i := range ids {
want[i] = ids[len(ids)-1-i]
}
for _, sort := range []string{"time", "severity"} {
t.Run(sort, func(t *testing.T) {
var got []int64
for page := 1; page <= 3; page++ {
items, total, pageErr := d.ListFindingsPage(FindingFilter{VulnClass: vc, Sort: sort}, page, 3)
if pageErr != nil {
t.Fatal(pageErr)
}
if total != len(ids) {
t.Fatalf("page %d total=%d, want %d", page, total, len(ids))
}
for _, item := range items {
got = append(got, item.ID)
}
}
if !slices.Equal(got, want) {
t.Fatalf("same-timestamp pagination was unstable: got=%v want=%v", got, want)
}
})
}
}
func TestFindingGroupsAndUnassignedPaging(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_groups_%d__", time.Now().UnixNano())
cleanupFindings := func() { _, _ = d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc) }
defer cleanupFindings()
taskA, err := d.CreateTask("group A", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskA.ID)
taskB, err := d.CreateTask("group B", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskB.ID)
seed := []struct {
taskID int64
severity string
status string
}{
{taskA.ID, SeverityCritical, FindingPending},
{taskA.ID, SeverityHigh, FindingResolved},
{taskB.ID, SeverityLow, FindingPending},
{0, SeverityMedium, FindingPending},
}
for i, item := range seed {
id, addErr := d.AddFinding(item.taskID, 0, vc, "", item.severity, fmt.Sprintf("summary %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
if _, setErr := d.SetFindingStatus(id, item.status); setErr != nil {
t.Fatalf("SetFindingStatus[%d]: %v", i, setErr)
}
}
// Retained findings from deleted tasks join the same bucket as findings that
// were created without any task.
if err := d.DeleteTask(taskB.ID); err != nil {
t.Fatal(err)
}
page, groupTotal, findingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc, Sort: "severity"}, 1, 1)
if err != nil {
t.Fatal(err)
}
if len(page) != 1 || groupTotal != 2 || findingTotal != 4 {
t.Fatalf("page totals: items=%d groups=%d findings=%d", len(page), groupTotal, findingTotal)
}
groups, _, _, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
var live, unassigned *FindingGroup
for i := range groups {
if groups[i].TaskID == nil {
unassigned = &groups[i]
} else if *groups[i].TaskID == taskA.ID {
live = &groups[i]
}
}
if live == nil || live.Count != 2 || live.Critical != 1 || live.High != 1 || live.TaskDescription != "group A" {
t.Fatalf("live group mismatch: %+v", live)
}
if unassigned == nil || unassigned.Count != 2 || unassigned.Medium != 1 || unassigned.Low != 1 {
t.Fatalf("unassigned group mismatch: %+v", unassigned)
}
groupForTask := func(items []FindingGroup, taskID int64) *FindingGroup {
for i := range items {
if items[i].TaskID != nil && *items[i].TaskID == taskID {
return &items[i]
}
}
return nil
}
if err := d.SetPaused(taskA.ID, true); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "paused" {
t.Fatalf("paused task group status mismatch: %+v", groups)
}
if err := d.SetPaused(taskA.ID, false); err != nil {
t.Fatal(err)
}
if err := d.Enqueue(taskA.ID, "resume"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "queued" {
t.Fatalf("queued task group status mismatch: %+v", groups)
}
if err := d.SetStatus(taskA.ID, "done"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "done" {
t.Fatalf("terminal task status must win over queued flag: %+v", groups)
}
orphans, total, err := d.ListFindingsPage(FindingFilter{VulnClass: vc, TaskID: FindingUnassignedTask}, 1, 10)
if err != nil || total != 2 || len(orphans) != 2 {
t.Fatalf("unassigned page: len=%d total=%d err=%v", len(orphans), total, err)
}
maxPage := int(^uint(0) >> 1)
farFindings, farTotal, err := d.ListFindingsPage(FindingFilter{VulnClass: vc}, maxPage, 200)
if err != nil || farTotal != len(seed) || len(farFindings) != 0 {
t.Fatalf("far finding page: len=%d total=%d err=%v", len(farFindings), farTotal, err)
}
farGroups, farGroupTotal, farFindingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, maxPage, 100)
if err != nil || farGroupTotal != 2 || farFindingTotal != len(seed) || len(farGroups) != 0 {
t.Fatalf("far group page: len=%d groups=%d findings=%d err=%v",
len(farGroups), farGroupTotal, farFindingTotal, err)
}
filtered, filteredGroups, filteredFindings, err := d.ListFindingGroups(
FindingFilter{VulnClass: vc, Status: FindingResolved}, 1, 10,
)
if err != nil || filteredGroups != 1 || filteredFindings != 1 || len(filtered) != 1 || filtered[0].High != 1 {
t.Fatalf("filtered groups: %+v groups=%d findings=%d err=%v", filtered, filteredGroups, filteredFindings, err)
}
}
func TestAddFindingFollowUpIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
task, err := d.CreateTask("finding follow-up", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("follow-up-%d.example.test", time.Now().UnixNano()),
TaskID: task.ID,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id=$1`, assetID)
store := d.Exploration(task.ExplorationID)
findingNodeID, err := store.AddNode(KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
findingID, err := d.AddFinding(task.ID, findingNodeID, "test", "source", SeverityHigh, "source", "", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
auditInput := Activity{Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "验证可利用性并形成证据链"}
intentID, audit, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "验证可利用性并形成证据链", auditInput)
if err != nil {
t.Fatal(err)
}
secondID, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "从另一条路径深入", Activity{
Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "从另一条路径深入",
})
if err != nil {
t.Fatal(err)
}
if secondID == intentID {
t.Fatal("repeated follow-up submissions must create distinct intents")
}
expectedAuditSummary := fmt.Sprintf("%s #%d", auditInput.Summary, intentID)
if audit.ID <= 0 || audit.NodeID == nil || *audit.NodeID != intentID || audit.CreatedAt.IsZero() || audit.Summary != expectedAuditSummary {
t.Fatalf("persisted audit mismatch: %+v", audit)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil {
t.Fatalf("GetNode: node=%+v err=%v", node, err)
}
if node.Kind != KindIntent || node.Priority != 10 || node.State != "open" || node.Origin != "human" {
t.Fatalf("follow-up intent metadata: %+v", node)
}
var anchorCount, edgeCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1 AND asset_id=$2`, intentID, assetID).Scan(&anchorCount); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_edges
WHERE exploration_id=$1 AND src_id=$2 AND rel=$3 AND dst_id=$4`,
task.ExplorationID, findingNodeID, RelDerivedFrom, intentID).Scan(&edgeCount); err != nil {
t.Fatal(err)
}
if anchorCount != 1 || edgeCount != 1 {
t.Fatalf("follow-up lineage: anchors=%d edges=%d", anchorCount, edgeCount)
}
var persistedAuditCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM activity
WHERE id=$1 AND exploration_id=$2 AND node_id=$3 AND worker='system' AND summary=$4`,
audit.ID, task.ExplorationID, intentID, expectedAuditSummary).Scan(&persistedAuditCount); err != nil {
t.Fatal(err)
}
if persistedAuditCount != 1 {
t.Fatalf("atomic audit count=%d, want 1", persistedAuditCount)
}
var nodesBefore, activitiesBefore int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesBefore); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesBefore); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "must roll back", Activity{
Worker: "system", Kind: "text", Summary: "invalid audit", Metadata: json.RawMessage(`{`),
}); err == nil {
t.Fatal("invalid activity metadata should fail the transaction")
}
var nodesAfter, activitiesAfter int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesAfter); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesAfter); err != nil {
t.Fatal(err)
}
if nodesAfter != nodesBefore || activitiesAfter != activitiesBefore {
t.Fatalf("activity failure did not roll back graph: nodes %d->%d activities %d->%d",
nodesBefore, nodesAfter, activitiesBefore, activitiesAfter)
}
if _, err := d.Exec(`DELETE FROM exploration_nodes WHERE id=$1`, findingNodeID); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "should fail", auditInput); !errors.Is(err, ErrFindingOriginUnavailable) {
t.Fatalf("missing source node: want ErrFindingOriginUnavailable, got %v", err)
}
}
db/intent_admission.go
+32
View File
@@ -0,0 +1,32 @@
package db
import "fmt"
// DiscardOpenIntent compensates a follow-up creation when task admission fails.
// The task execution gate must still be held by the caller, so the intent cannot
// be claimed between this check and deletion. Edges and anchors cascade with the
// node; activity is deleted explicitly because its node FK otherwise becomes NULL.
func (s *ExplorationStore) DiscardOpenIntent(id int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil {
return err
}
res, err := tx.Exec(`DELETE FROM exploration_nodes
WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID)
if err != nil {
return err
}
removed, err := res.RowsAffected()
if err != nil {
return err
}
if removed != 1 {
return fmt.Errorf("open intent %d was not available for admission rollback", id)
}
return tx.Commit()
}
db/intent_control_test.go
+285
View File
@@ -0,0 +1,285 @@
package db
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("intent CAS", "only one controller wins")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
const contenders = 12
var winners atomic.Int32
var wg sync.WaitGroup
start := make(chan struct{})
for range contenders {
wg.Add(1)
go func() {
defer wg.Done()
<-start
changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
if transitionErr != nil {
t.Errorf("transition: %v", transitionErr)
return
}
if changed {
winners.Add(1)
}
}()
}
close(start)
wg.Wait()
if got := winners.Load(); got != 1 {
t.Fatalf("CAS winners=%d, want 1", got)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil || node.State != "paused" {
t.Fatalf("node=%+v err=%v, want paused", node, err)
}
}
func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
baselineDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
baseline := dailyTokenBuckets(baselineDaily)
appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
t.Helper()
activityID, appendErr := store.AppendActivity(Activity{
NodeID: &intentID, Worker: "work#1", Kind: kind,
InputTokens: input, OutputTokens: output,
CacheReadTokens: read, CacheWriteTokens: write,
})
if appendErr != nil {
t.Fatal(appendErr)
}
if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
t.Fatal(updateErr)
}
}
values := func(input, output, read, write int) (*int, *int, *int, *int) {
return &input, &output, &read, &write
}
now := time.Now().UTC()
atUTCNoon := func(daysAgo int) time.Time {
at := now.AddDate(0, 0, -daysAgo)
return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
}
authoritativeDay := atUTCNoon(6)
fallbackDay := atUTCNoon(4)
trailingDay := atUTCNoon(2)
// Run 1 has an authoritative result; its preceding cumulative frame must not
// be counted a second time or move usage to the frame's earlier date.
i, o, r, w := values(100, 100, 100, 100)
appendUsage("usage", atUTCNoon(7), i, o, r, w)
i, o, r, w = values(10, 11, 12, 13)
appendUsage("result", authoritativeDay, i, o, r, w)
// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
// to the result's date, not the cumulative frame's date.
i, o, r, w = values(20, 21, 22, 23)
appendUsage("usage", atUTCNoon(5), i, o, r, w)
appendUsage("result", fallbackDay, nil, nil, nil, nil)
// Run 3 was interrupted before a terminal result was persisted.
i, o, r, w = values(30, 31, 32, 33)
appendUsage("usage", trailingDay, i, o, r, w)
beforeCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
beforeCancel := dailyTokenBuckets(beforeCancelDaily)
assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
if err := store.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(intentID)
if err != nil {
t.Fatal(err)
}
if cleanup.Activities != 5 {
t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total, 60, 63, 66, 69)
afterCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
afterCancel := dailyTokenBuckets(afterCancelDaily)
assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
if len(sessions) != 0 {
t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
}
var rollups, datedRollups int
if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
) FROM activity
WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
t.Fatal(err)
}
if rollups != 3 || datedRollups != rollups {
t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
}
if _, err := store.CancelIntent(intentID); err == nil {
t.Fatal("second cancellation unexpectedly succeeded")
}
afterRetry, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
afterRetryDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
}
func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
out := make(map[string]DailyTokenBucket, len(items))
for _, item := range items {
out[item.Day] = item
}
return out
}
func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
t.Helper()
day := at.UTC().Format(time.DateOnly)
actual, initial := got[day], baseline[day]
if actual.InputTokens-initial.InputTokens != input ||
actual.OutputTokens-initial.OutputTokens != output ||
actual.CacheReadTokens-initial.CacheReadTokens != read {
t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
}
}
func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
}
for day, expected := range want {
if actual, ok := got[day]; !ok || actual != expected {
t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
}
}
}
func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
for _, intentID := range []int64{first, second} {
if err := store.Link(intentID, RelYields, fact); err != nil {
t.Fatal(err)
}
if err := store.Link(intentID, RelYields, finding); err != nil {
t.Fatal(err)
}
}
if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
if err := store.SetIntentState(first, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(first)
if err != nil {
t.Fatal(err)
}
if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
}
for _, nodeID := range []int64{fact, finding} {
node, getErr := store.GetNode(nodeID)
if getErr != nil || node == nil {
t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
}
}
}
db/schema.sql
+59 -2
View File
@@ -31,6 +31,11 @@ CREATE TABLE IF NOT EXISTS assets (
'root_domain','ip','subdomain','app','service','endpoint'
)),
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
-- explicit: caller/user selected the company; scope: derived from company_scope.
-- Existing installations are conservatively migrated as explicit so a scope
-- rebuild can never erase a historical manual association.
company_source TEXT NOT NULL DEFAULT 'explicit'
CHECK (company_source IN ('explicit','scope')),
task_ids BIGINT[] NOT NULL DEFAULT '{}',
domain TEXT,
root_domain TEXT,
@@ -84,6 +89,13 @@ CREATE INDEX IF NOT EXISTS idx_av2_bound_domains ON assets USING GIN(bound_domai
CREATE INDEX IF NOT EXISTS idx_av2_open_ports ON assets USING GIN(open_ports) WHERE type = 'ip';
CREATE INDEX IF NOT EXISTS idx_av2_last_seen ON assets(last_seen DESC);
CREATE INDEX IF NOT EXISTS idx_av2_type_seen ON assets(type, last_seen DESC);
ALTER TABLE assets ADD COLUMN IF NOT EXISTS company_source TEXT;
UPDATE assets SET company_source = 'explicit' WHERE company_source IS NULL;
ALTER TABLE assets ALTER COLUMN company_source SET DEFAULT 'explicit';
ALTER TABLE assets ALTER COLUMN company_source SET NOT NULL;
ALTER TABLE assets DROP CONSTRAINT IF EXISTS assets_company_source_check;
ALTER TABLE assets ADD CONSTRAINT assets_company_source_check
CHECK (company_source IN ('explicit','scope'));
DROP TRIGGER IF EXISTS trg_av2_upd ON assets;
CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
@@ -91,17 +103,36 @@ CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
CREATE TABLE IF NOT EXISTS company_scope (
id BIGSERIAL PRIMARY KEY,
company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr')),
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr','icp','keyword')),
domain TEXT,
net CIDR,
value TEXT,
raw TEXT NOT NULL,
reason TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT uq_sv2_domain UNIQUE (company_id, domain),
CONSTRAINT uq_sv2_net UNIQUE (company_id, net)
CONSTRAINT uq_sv2_net UNIQUE (company_id, net),
CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
)
);
-- Existing installations need the new text payload and expanded kind check.
ALTER TABLE company_scope ADD COLUMN IF NOT EXISTS value TEXT;
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS company_scope_kind_check;
ALTER TABLE company_scope ADD CONSTRAINT company_scope_kind_check
CHECK (kind IN ('domain','ip','cidr','icp','keyword'));
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS ck_company_scope_payload;
ALTER TABLE company_scope ADD CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
);
CREATE INDEX IF NOT EXISTS idx_sv2_domain ON company_scope(domain) WHERE kind = 'domain';
CREATE INDEX IF NOT EXISTS idx_sv2_net ON company_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr');
CREATE UNIQUE INDEX IF NOT EXISTS uq_sv2_value ON company_scope(company_id, kind, value) WHERE kind IN ('icp','keyword');
CREATE INDEX IF NOT EXISTS idx_sv2_icp ON company_scope(value) WHERE kind = 'icp';
CREATE INDEX IF NOT EXISTS idx_sv2_company ON company_scope(company_id);
-- =====================================================================
@@ -301,6 +332,8 @@ CREATE TABLE IF NOT EXISTS tasks (
CHECK (status IN ('created','running','paused','done','failed','timeout')),
paused BOOLEAN NOT NULL DEFAULT false,
queued BOOLEAN NOT NULL DEFAULT false,
queued_at TIMESTAMPTZ,
queue_mode TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
llm_chain_revision BIGINT NOT NULL DEFAULT 0,
@@ -324,9 +357,30 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300;
-- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false;
-- queued_at makes admission FIFO reflect the actual enqueue order rather than the
-- task creation order. queue_mode distinguishes first bootstrap from resuming an
-- exploration that already owns goals/history.
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued_at TIMESTAMPTZ;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queue_mode TEXT NOT NULL DEFAULT '';
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_chain_revision BIGINT NOT NULL DEFAULT 0;
-- Reusable task description/goal presets. nkey is the normalized, case-insensitive
-- identity used to reject visually equivalent duplicate names.
CREATE TABLE IF NOT EXISTS task_templates (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
nkey TEXT NOT NULL UNIQUE,
description TEXT NOT NULL,
goal TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_task_templates_updated ON task_templates(updated_at DESC, id DESC);
DROP TRIGGER IF EXISTS trg_task_templates_upd ON task_templates;
CREATE TRIGGER trg_task_templates_upd BEFORE UPDATE ON task_templates
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- Direct, read-only task context inheritance. Relations are intentionally not
-- recursive: a task sees only the source tasks explicitly chosen at creation.
CREATE TABLE IF NOT EXISTS task_relations (
@@ -573,10 +627,13 @@ CREATE TABLE IF NOT EXISTS conversations (
agent_key TEXT NOT NULL,
title TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
pinned_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE conversations ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ;
CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned_at DESC) WHERE pinned_at IS NOT NULL;
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
db/task_context.go
+33 -2
View File
@@ -60,6 +60,11 @@ func (d *DB) hydrateTaskContext(t *Task) error {
return err
}
t.SourceTaskIDs = sources
companies, err := d.TaskCompanyIDs(t.ID)
if err != nil {
return err
}
t.CompanyIDs = companies
t.LLMProfileIDs = make([]int64, 0, len(chain))
t.LLMFailoverState = "default"
var latest *TaskLLMProfile
@@ -172,6 +177,28 @@ func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) {
return out, rows.Err()
}
// TaskCompanyIDs returns the companies whose asset scopes are available to the
// task. The task_scope rows remain the single source of truth.
func (d *DB) TaskCompanyIDs(taskID int64) ([]int64, error) {
rows, err := d.Query(`
SELECT company_id FROM task_scope
WHERE task_id=$1 AND kind='company' AND company_id IS NOT NULL
ORDER BY id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) {
rows, err := d.Query(`
SELECT t.id, t.exploration_id, t.description, t.goal, t.status
@@ -339,13 +366,17 @@ ORDER BY position LIMIT 1`, taskID, position).Scan(&next)
case nil:
out.Advanced = true
out.NextProfileID = &next
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=$2, llm_profile_id=$2 WHERE id=$1`, taskID, next); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID, next); err != nil {
return out, err
}
case sql.ErrNoRows:
out.Advanced = true
out.ChainExhausted = true
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=NULL, llm_profile_id=NULL WHERE id=$1`, taskID); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID); err != nil {
return out, err
}
default:
db/task_queue_test.go
+66
View File
@@ -0,0 +1,66 @@
package db
import (
"testing"
"time"
)
func TestTaskQueuePreservesBootstrapAndFIFOPosition(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("queue metadata", "keep first-run mode", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
first, err := d.GetTask(task.ID)
if err != nil || first == nil || first.QueuedAt == nil {
t.Fatalf("first enqueue: task=%+v err=%v", first, err)
}
// A follow-up/rerun can try to admit an already queued task as resume. The
// original bootstrap mode and FIFO timestamp must remain authoritative.
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "resume"); err != nil {
t.Fatal(err)
}
second, err := d.GetTask(task.ID)
if err != nil || second == nil {
t.Fatalf("second enqueue: task=%+v err=%v", second, err)
}
if second.QueueMode != "bootstrap" {
t.Fatalf("queue mode=%q, want bootstrap", second.QueueMode)
}
if second.QueuedAt == nil || !second.QueuedAt.Equal(*first.QueuedAt) {
t.Fatalf("repeated enqueue moved FIFO position: first=%v second=%v", first.QueuedAt, second.QueuedAt)
}
// Pausing a queued task removes it from the queue but keeps its required
// startup mode. A later requeue receives a new tail position.
if err := d.Dequeue(task.ID, false); err != nil {
t.Fatal(err)
}
paused, err := d.GetTask(task.ID)
if err != nil || paused == nil || paused.Queued || paused.QueueMode != "bootstrap" {
t.Fatalf("paused queue metadata: task=%+v err=%v", paused, err)
}
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
requeued, err := d.GetTask(task.ID)
if err != nil || requeued == nil || requeued.QueuedAt == nil {
t.Fatalf("requeue: task=%+v err=%v", requeued, err)
}
if !requeued.QueuedAt.After(*first.QueuedAt) {
t.Fatalf("requeue did not move to FIFO tail: first=%v requeued=%v", first.QueuedAt, requeued.QueuedAt)
}
}
db/task_source_limit_test.go
+27
View File
@@ -1,6 +1,8 @@
package db
import (
"errors"
"reflect"
"strings"
"testing"
)
@@ -18,3 +20,28 @@ func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) {
t.Fatalf("expected source-count validation error, got %v", err)
}
}
func TestNormalizeTaskCompanyIDs(t *testing.T) {
got, err := NormalizeTaskCompanyIDs([]int64{4, 2, 4, 7, 2})
if err != nil {
t.Fatal(err)
}
if want := []int64{4, 2, 7}; !reflect.DeepEqual(got, want) {
t.Fatalf("NormalizeTaskCompanyIDs=%v, want %v", got, want)
}
if _, err := NormalizeTaskCompanyIDs([]int64{1, 0}); !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("invalid company id error=%v", err)
}
}
func TestCreateTaskRejectsTooManyCompaniesBeforeOpeningTransaction(t *testing.T) {
companyIDs := make([]int64, MaxTaskCompanyCount+1)
for i := range companyIDs {
companyIDs[i] = int64(i + 1)
}
_, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{CompanyIDs: companyIDs})
if !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("expected company-count validation error, got %v", err)
}
}
db/task_templates.go
+231
View File
@@ -0,0 +1,231 @@
package db
import (
"database/sql"
"errors"
"fmt"
"strings"
"time"
"unicode/utf8"
"github.com/jackc/pgx/v5/pgconn"
)
const (
MaxTaskTemplateNameRunes = 120
MaxTaskTemplateTextRunes = 16000
)
var (
ErrTaskTemplateInvalid = errors.New("invalid task template")
ErrTaskTemplateNameConflict = errors.New("task template name already exists")
ErrTaskTemplateNotFound = errors.New("task template not found")
)
// TaskTemplate is a reusable task description and goal preset.
type TaskTemplate struct {
ID int64 `json:"id"`
Name string `json:"name"`
NKey string `json:"-"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// TaskTemplateInput is the create/update payload after normalization.
type TaskTemplateInput struct {
Name string
Description string
Goal string
}
// TaskTemplatePatch changes only fields whose pointers are non-nil.
type TaskTemplatePatch struct {
Name *string
Description *string
Goal *string
}
const taskTemplateCols = `id, name, nkey, description, goal, created_at, updated_at`
func scanTaskTemplate(row interface{ Scan(...any) error }) (TaskTemplate, error) {
var t TaskTemplate
err := row.Scan(&t.ID, &t.Name, &t.NKey, &t.Description, &t.Goal, &t.CreatedAt, &t.UpdatedAt)
return t, err
}
// taskTemplateName normalizes display whitespace while preserving the user's case.
func taskTemplateName(name string) string { return strings.Join(strings.Fields(name), " ") }
// taskTemplateNKey is the case-insensitive identity used by the unique index.
func taskTemplateNKey(name string) string { return strings.ToLower(taskTemplateName(name)) }
func normalizeTaskTemplateInput(in TaskTemplateInput) (TaskTemplateInput, string, error) {
in.Name = taskTemplateName(in.Name)
in.Description = strings.TrimSpace(in.Description)
in.Goal = strings.TrimSpace(in.Goal)
switch {
case in.Name == "":
return in, "", fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Name) > MaxTaskTemplateNameRunes:
return in, "", fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
case in.Description == "":
return in, "", fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Description) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
case in.Goal == "":
return in, "", fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Goal) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
return in, taskTemplateNKey(in.Name), nil
}
func normalizeTaskTemplatePatch(patch TaskTemplatePatch) (TaskTemplatePatch, *string, error) {
if patch.Name == nil && patch.Description == nil && patch.Goal == nil {
return patch, nil, fmt.Errorf("%w: no fields supplied", ErrTaskTemplateInvalid)
}
var nkey *string
if patch.Name != nil {
name := taskTemplateName(*patch.Name)
if name == "" {
return patch, nil, fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(name) > MaxTaskTemplateNameRunes {
return patch, nil, fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
}
key := taskTemplateNKey(name)
patch.Name = &name
nkey = &key
}
if patch.Description != nil {
description := strings.TrimSpace(*patch.Description)
if description == "" {
return patch, nil, fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(description) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Description = &description
}
if patch.Goal != nil {
goal := strings.TrimSpace(*patch.Goal)
if goal == "" {
return patch, nil, fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(goal) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Goal = &goal
}
return patch, nkey, nil
}
func taskTemplateUniqueViolation(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "23505"
}
// CreateTaskTemplate inserts one globally reusable preset.
func (d *DB) CreateTaskTemplate(in TaskTemplateInput) (*TaskTemplate, error) {
in, nkey, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`
INSERT INTO task_templates(name, nkey, description, goal)
VALUES ($1,$2,$3,$4)
ON CONFLICT (nkey) DO NOTHING
RETURNING `+taskTemplateCols, in.Name, nkey, in.Description, in.Goal))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// ListTaskTemplates returns the most recently maintained templates first.
func (d *DB) ListTaskTemplates() ([]*TaskTemplate, error) {
rows, err := d.Query(`SELECT ` + taskTemplateCols + ` FROM task_templates ORDER BY updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*TaskTemplate{}
for rows.Next() {
t, err := scanTaskTemplate(rows)
if err != nil {
return nil, err
}
out = append(out, &t)
}
return out, rows.Err()
}
// GetTaskTemplate returns nil when id does not exist.
func (d *DB) GetTaskTemplate(id int64) (*TaskTemplate, error) {
t, err := scanTaskTemplate(d.QueryRow(`SELECT `+taskTemplateCols+` FROM task_templates WHERE id=$1`, id))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &t, nil
}
// UpdateTaskTemplate replaces the editable fields of one preset.
func (d *DB) UpdateTaskTemplate(id int64, in TaskTemplateInput) (*TaskTemplate, error) {
in, _, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
return d.PatchTaskTemplate(id, TaskTemplatePatch{
Name: &in.Name, Description: &in.Description, Goal: &in.Goal,
})
}
// PatchTaskTemplate atomically changes only the supplied fields. Keeping the
// merge in one UPDATE prevents concurrent disjoint PATCH requests from losing
// each other's changes.
func (d *DB) PatchTaskTemplate(id int64, patch TaskTemplatePatch) (*TaskTemplate, error) {
patch, nkey, err := normalizeTaskTemplatePatch(patch)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`UPDATE task_templates
SET name=CASE WHEN $2 THEN $3::text ELSE name END,
nkey=CASE WHEN $2 THEN $4::text ELSE nkey END,
description=CASE WHEN $5 THEN $6::text ELSE description END,
goal=CASE WHEN $7 THEN $8::text ELSE goal END
WHERE id=$1
RETURNING `+taskTemplateCols,
id,
patch.Name != nil, patch.Name, nkey,
patch.Description != nil, patch.Description,
patch.Goal != nil, patch.Goal,
))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNotFound
}
if taskTemplateUniqueViolation(err) {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// DeleteTaskTemplate deletes one preset and reports whether it existed.
func (d *DB) DeleteTaskTemplate(id int64) (bool, error) {
result, err := d.Exec(`DELETE FROM task_templates WHERE id=$1`, id)
if err != nil {
return false, err
}
n, err := result.RowsAffected()
return n > 0, err
}
db/task_templates_test.go
+131
View File
@@ -0,0 +1,131 @@
package db
import (
"errors"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestTaskTemplateCRUDAndNormalizedUniqueness(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
name := fmt.Sprintf("Template %d", suffix)
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: " " + strings.ReplaceAll(name, " ", " ") + " ",
Description: " initial description ",
Goal: " initial goal ",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
if created.Name != name || created.Description != "initial description" || created.Goal != "initial goal" {
t.Fatalf("template was not normalized: %+v", created)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: strings.ToUpper(name), Description: "duplicate", Goal: "duplicate",
}); !errors.Is(err, ErrTaskTemplateNameConflict) {
t.Fatalf("duplicate create error = %v, want %v", err, ErrTaskTemplateNameConflict)
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil || got == nil || got.Name != name {
t.Fatalf("GetTaskTemplate = %+v, %v", got, err)
}
listed, err := d.ListTaskTemplates()
if err != nil {
t.Fatal(err)
}
found := false
for _, template := range listed {
if template.ID == created.ID {
found = true
break
}
}
if !found {
t.Fatalf("created template %d missing from list", created.ID)
}
updatedName := name + " updated"
updated, err := d.UpdateTaskTemplate(created.ID, TaskTemplateInput{
Name: updatedName, Description: "new description", Goal: "new goal",
})
if err != nil {
t.Fatal(err)
}
if updated.Name != updatedName || updated.Description != "new description" || updated.Goal != "new goal" {
t.Fatalf("unexpected updated template: %+v", updated)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{Name: "", Description: "x", Goal: "y"}); !errors.Is(err, ErrTaskTemplateInvalid) {
t.Fatalf("empty name error = %v, want %v", err, ErrTaskTemplateInvalid)
}
deleted, err := d.DeleteTaskTemplate(created.ID)
if err != nil || !deleted {
t.Fatalf("DeleteTaskTemplate = %v, %v", deleted, err)
}
deleted, err = d.DeleteTaskTemplate(created.ID)
if err != nil || deleted {
t.Fatalf("second DeleteTaskTemplate = %v, %v", deleted, err)
}
}
func TestTaskTemplateDisjointPatchesCompose(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: fmt.Sprintf("Concurrent template %d", time.Now().UnixNano()),
Description: "initial description",
Goal: "initial goal",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
description := "description from concurrent patch"
goal := "goal from concurrent patch"
start := make(chan struct{})
errs := make(chan error, 2)
var wg sync.WaitGroup
for _, patch := range []TaskTemplatePatch{{Description: &description}, {Goal: &goal}} {
patch := patch
wg.Add(1)
go func() {
defer wg.Done()
<-start
_, err := d.PatchTaskTemplate(created.ID, patch)
errs <- err
}()
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatal(err)
}
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil {
t.Fatal(err)
}
if got == nil || got.Description != description || got.Goal != goal {
t.Fatalf("disjoint patches lost an update: %+v", got)
}
}
db/tasks.go
+119 -12
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"time"
@@ -10,14 +11,16 @@ import (
// Task is a row in the task registry (1:1 with an exploration).
type Task struct {
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
QueuedAt *time.Time `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
// for ActiveLLMProfileID while older API clients still send one profile id.
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
@@ -26,6 +29,7 @@ type Task struct {
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
@@ -74,6 +78,40 @@ const MinPlanHeartbeatSeconds = 600
// multiplying graph and asset-context queries without limit.
const MaxTaskSourceCount = 8
// MaxTaskCompanyCount bounds the number of company asset scopes attached to a
// task. Company scopes are prompt context, so an unbounded list would inflate
// every planner and worker call.
const MaxTaskCompanyCount = 32
var (
ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids")
ErrTaskCompanyNotFound = errors.New("task company not found")
)
// NormalizeTaskCompanyIDs validates IDs and removes duplicates while retaining
// the user's first-seen order.
func NormalizeTaskCompanyIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, nil
}
seen := make(map[int64]struct{}, min(len(ids), MaxTaskCompanyCount))
normalized := make([]int64, 0, min(len(ids), MaxTaskCompanyCount))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("%w: company id must be positive", ErrTaskCompanyIDsInvalid)
}
if _, exists := seen[id]; exists {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
if len(normalized) > MaxTaskCompanyCount {
return nil, fmt.Errorf("%w: got more than %d unique companies", ErrTaskCompanyIDsInvalid, MaxTaskCompanyCount)
}
}
return normalized, nil
}
func normalizeHeartbeat(sec int) int {
if sec < MinPlanHeartbeatSeconds {
return MinPlanHeartbeatSeconds
@@ -95,6 +133,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe
// with the task/exploration row.
type TaskCreateOptions struct {
SourceTaskIDs []int64
CompanyIDs []int64
LLMProfileIDs []int64
TimeoutSeconds int
PlanHeartbeatSeconds int
@@ -106,6 +145,11 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti
if len(opts.SourceTaskIDs) > MaxTaskSourceCount {
return nil, fmt.Errorf("too many source tasks: got %d, maximum is %d", len(opts.SourceTaskIDs), MaxTaskSourceCount)
}
companyIDs, err := NormalizeTaskCompanyIDs(opts.CompanyIDs)
if err != nil {
return nil, err
}
opts.CompanyIDs = companyIDs
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -144,6 +188,7 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload));
LLMProfileID: active, ActiveLLMProfileID: active,
LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...),
SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), opts.CompanyIDs...),
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
}
if err := tx.QueryRow(`
@@ -155,6 +200,9 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil {
return nil, err
}
if err := insertTaskCompanies(tx, t.ID, opts.CompanyIDs); err != nil {
return nil, err
}
if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil {
return nil, err
}
@@ -166,6 +214,33 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
return t, tx.Commit()
}
func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
if len(companyIDs) == 0 {
return nil
}
var inserted int
err := tx.QueryRow(`
WITH requested(company_id, position) AS (
SELECT company_id, position
FROM unnest($2::bigint[]) WITH ORDINALITY AS requested(company_id, position)
), inserted AS (
INSERT INTO task_scope(task_id, kind, company_id, source, reason)
SELECT $1, 'company', companies.id, 'manual', '任务创建时关联企业'
FROM requested
JOIN companies ON companies.id=requested.company_id
ORDER BY requested.position
RETURNING company_id
)
SELECT count(*) FROM inserted`, taskID, companyIDs).Scan(&inserted)
if err != nil {
return err
}
if inserted != len(companyIDs) {
return fmt.Errorf("%w: one or more companies do not exist", ErrTaskCompanyNotFound)
}
return nil
}
func insertTaskRelations(tx *sql.Tx, taskID int64, sourceIDs []int64) error {
seen := map[int64]bool{}
for _, sourceID := range sourceIDs {
@@ -199,11 +274,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
return nil
}
const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
const taskCols = `id, description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
return nil, err
}
return &t, nil
@@ -263,11 +338,43 @@ func (d *DB) SetPaused(id int64, paused bool) error {
return err
}
func (d *DB) SetQueued(id int64, queued bool) error {
_, err := d.Exec(`UPDATE tasks SET queued=$1 WHERE id=$2`, queued, id)
// Enqueue places a task at the tail of the persistent admission queue. Repeating
// the operation while it is already queued keeps its original FIFO position.
func (d *DB) Enqueue(id int64, mode string) error {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
_, err := d.Exec(`UPDATE tasks
SET queued=true,
queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END,
queue_mode=CASE
WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END
WHERE id=$1`, id, mode)
return err
}
// Dequeue removes the concurrency hold. clearMode is false when a user pauses a
// queued task, preserving whether its next admission must bootstrap or resume.
func (d *DB) Dequeue(id int64, clearMode bool) error {
_, err := d.Exec(`UPDATE tasks
SET queued=false,
queued_at=NULL,
queue_mode=CASE WHEN $2 THEN '' ELSE queue_mode END
WHERE id=$1`, id, clearMode)
return err
}
// SetQueued is the compatibility helper used by older callers and tests. New
// scheduling code should use Enqueue/Dequeue so FIFO metadata is explicit.
func (d *DB) SetQueued(id int64, queued bool) error {
if queued {
return d.Enqueue(id, "bootstrap")
}
return d.Dequeue(id, true)
}
// SetStatus updates a task's lifecycle status. Entering a terminal state
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
db/tasks_test.go
+80
View File
@@ -1,6 +1,7 @@
package db
import (
"errors"
"fmt"
"sync"
"testing"
@@ -277,6 +278,7 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] {
t.Fatalf("replacement revision/cursor not preserved: %+v", got)
}
replacementRevision := got.LLMChainRevision
transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota")
if err != nil {
@@ -285,6 +287,13 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] {
t.Fatalf("unexpected first transition: %+v", transition)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.LLMChainRevision != replacementRevision+1 {
t.Fatalf("automatic cursor advance did not increment revision: before=%d after=%d", replacementRevision, got.LLMChainRevision)
}
late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate")
if err != nil {
t.Fatal(err)
@@ -483,3 +492,74 @@ func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) {
t.Fatal("terminal task LLM edit should be rejected")
}
}
func TestCreateTaskWithCompanyScopes(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "")
if err != nil {
t.Fatal(err)
}
companyB, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company B %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB)
})
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyB},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("creation result company IDs=%v", task.CompanyIDs)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("hydrated company IDs=%v", got.CompanyIDs)
}
var scopeCount int
if err := d.QueryRow(`SELECT count(*) FROM task_scope WHERE task_id=$1 AND kind='company'`, task.ID).Scan(&scopeCount); err != nil {
t.Fatal(err)
}
if scopeCount != 2 {
t.Fatalf("company task scope rows=%d want 2", scopeCount)
}
duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyA, companyB, companyA},
})
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(duplicateTask.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("company IDs were not normalized: %v", duplicateTask.CompanyIDs)
}
if err := d.DeleteTask(duplicateTask.ID); err != nil {
t.Fatal(err)
}
badDescription := fmt.Sprintf("invalid-company-%d", suffix)
if _, err := d.CreateTaskWithOptions(badDescription, "rollback", TaskCreateOptions{CompanyIDs: []int64{companyA, 1 << 62}}); !errors.Is(err, ErrTaskCompanyNotFound) {
t.Fatalf("missing company error=%v, want %v", err, ErrTaskCompanyNotFound)
}
var leaked int
if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, badDescription).Scan(&leaked); err != nil {
t.Fatal(err)
}
if leaked != 0 {
t.Fatalf("failed company association leaked %d exploration rows", leaked)
}
}
server/assets.go
+151 -52
View File
@@ -2,12 +2,58 @@ package server
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/db"
)
// companyScopeInputs accepts both the new [{kind,value}] contract and the
// historical ["example.com","203.0.113.10"] contract.
type companyScopeInputs []db.ScopeInput
const maxCompanyMutationBodyBytes = 2 << 20
func decodeCompanyMutationRequest(w http.ResponseWriter, r *http.Request, value any) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxCompanyMutationBodyBytes)
if err := json.NewDecoder(r.Body).Decode(value); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return false
}
return true
}
func (items *companyScopeInputs) UnmarshalJSON(data []byte) error {
var rawItems []json.RawMessage
if err := json.Unmarshal(data, &rawItems); err != nil {
return err
}
out := make([]db.ScopeInput, 0, len(rawItems))
for i, raw := range rawItems {
var legacy string
if err := json.Unmarshal(raw, &legacy); err == nil {
out = append(out, db.ScopeInput{Value: legacy})
continue
}
var structured db.ScopeInput
if err := json.Unmarshal(raw, &structured); err != nil {
return fmt.Errorf("scope[%d] must be a string or {kind,value}", i)
}
out = append(out, structured)
}
*items = out
return nil
}
// assetStore returns the asset store.
func (s *Server) assetStore() *db.AssetStore {
if s.m.pg == nil {
@@ -53,32 +99,45 @@ func (s *Server) createCompany(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Name string `json:"name"`
Logo string `json:"logo"`
Scope []string `json:"scope"`
Name string `json:"name"`
Logo string `json:"logo"`
Scope companyScopeInputs `json:"scope"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
req.Name = strings.TrimSpace(req.Name)
if req.Name == "" {
writeErr(w, 400, "name required")
return
}
id, created, err := cs.UpsertCompany(req.Name, req.Logo)
if err != nil {
writeErr(w, 500, err.Error())
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
out := map[string]any{"id": id, "created": created}
if len(req.Scope) > 0 {
added, skipped, invalid, errs := cs.AddScope(id, req.Scope, "api")
out["scope_added"] = added
out["scope_skipped"] = skipped
out["scope_invalid"] = invalid
if len(errs) > 0 {
out["scope_errors"] = errs
id, added, skipped, invalid, scopeErrs, err := cs.CreateCompanyWithScope(req.Name, req.Logo, req.Scope, "api")
if err != nil {
if errors.Is(err, db.ErrCompanyNameConflict) {
writeErr(w, http.StatusConflict, "企业名称已存在")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(err, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
out := map[string]any{
"id": id,
"created": true,
"scope_added": added,
"scope_skipped": skipped,
"scope_invalid": invalid,
}
if len(scopeErrs) > 0 {
out["scope_errors"] = scopeErrs
}
writeJSON(w, 201, out)
}
@@ -131,20 +190,37 @@ func (s *Server) addCompanyScope(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Scope []string `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
Scope companyScopeInputs `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
var added, skipped, invalid int
var errs []string
var mutationErr error
if req.Reset {
added, invalid, errs = cs.UpdateScope(id, req.Scope, req.Reason)
added, invalid, errs, mutationErr = cs.UpdateScopeInputsChecked(id, req.Scope, req.Reason)
} else {
added, skipped, invalid, errs = cs.AddScope(id, req.Scope, req.Reason)
added, skipped, invalid, errs, mutationErr = cs.AddScopeInputsChecked(id, req.Scope, req.Reason)
}
if mutationErr != nil {
if errors.Is(mutationErr, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(mutationErr, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, mutationErr.Error())
return
}
out := map[string]any{
"added": added,
@@ -175,21 +251,32 @@ func (s *Server) deleteCompany(w http.ResponseWriter, r *http.Request) {
var req struct {
DeleteAssets bool `json:"delete_assets"`
}
// Body is optional; ignore decode errors (e.g. empty body).
_ = json.NewDecoder(r.Body).Decode(&req)
assetsDeleted := int64(0)
if req.DeleteAssets {
as := s.assetStore()
if as != nil {
assetsDeleted, err = as.DeleteByCompanyID(id)
if err != nil {
writeErr(w, 500, err.Error())
return
// The body is optional. Only an actual empty body is ignored; malformed or
// trailing JSON is a client error.
decoder := json.NewDecoder(r.Body)
if err := decoder.Decode(&req); err != nil {
if !errors.Is(err, io.EOF) {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
return
}
} else {
var trailing any
if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
if err == nil {
writeErr(w, http.StatusBadRequest, "invalid JSON: multiple values")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return
}
}
if err := cs.DeleteCompany(id); err != nil {
assetsDeleted, err := s.m.DeleteCompanyWithAssets(id, req.DeleteAssets)
if err != nil {
if errors.Is(err, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
writeErr(w, 500, err.Error())
return
}
@@ -217,7 +304,10 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
// GET /api/assets
// =====================================================================
const defaultAssetPageSize = 50
const (
defaultAssetPageSize = 50
maxAssetPageSize = 200
)
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
as := s.assetStore()
@@ -231,46 +321,55 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
offset, _ := strconv.Atoi(q.Get("offset"))
if limit <= 0 {
limit = defaultAssetPageSize
} else if limit > maxAssetPageSize {
limit = maxAssetPageSize
}
if offset < 0 {
offset = 0
}
var assets []*db.Asset
assets := []*db.Asset{}
var err error
// total is the full match count ignoring limit/offset — lets the UI paginate on
// the server. For company/task scopes it falls back to the returned length.
// total is the full match count ignoring limit/offset. Count first so an offset
// beyond the last row can return an empty page without an expensive scan.
total := 0
if dsl := q.Get("dsl"); dsl != "" {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
if err == nil {
total, _ = as.CountDSL(dsl, typ)
if err := db.ValidateDSL(dsl); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
total, err = as.CountDSL(dsl, typ)
if err == nil && offset < total {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
}
} else {
companyID, _ := strconv.ParseInt(q.Get("company_id"), 10, 64)
taskID, _ := strconv.ParseInt(q.Get("task_id"), 10, 64)
switch {
case companyID > 0:
assets, err = as.QueryByCompany(companyID, typ, limit, offset)
if err == nil {
total, err = as.CountByCompany(companyID, typ)
total, err = as.CountByCompany(companyID, typ)
if err == nil && offset < total {
assets, err = as.QueryByCompany(companyID, typ, limit, offset)
}
case taskID > 0:
assets, err = as.QueryByTask(taskID, typ, limit, offset)
if err == nil {
total, err = as.CountByTask(taskID, typ)
total, err = as.CountByTask(taskID, typ)
if err == nil && offset < total {
assets, err = as.QueryByTask(taskID, typ, limit, offset)
}
default:
if typ == "" {
typ = "root_domain"
}
assets, err = as.QueryByType(typ, limit, offset)
if err == nil {
total, _ = as.CountByType(typ)
total, err = as.CountByType(typ)
if err == nil && offset < total {
assets, err = as.QueryByType(typ, limit, offset)
}
}
}
if err != nil {
writeErr(w, 400, err.Error())
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, 200, map[string]any{
server/assets_scope_test.go
+265
View File
@@ -0,0 +1,265 @@
package server
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestCompanyScopeInputsAcceptStructuredAndLegacyRules(t *testing.T) {
var inputs companyScopeInputs
if err := json.Unmarshal([]byte(`[
"example.com",
{"kind":"icp","value":"京 ICP备 123号"},
{"kind":"keyword","value":"Acme Security"}
]`), &inputs); err != nil {
t.Fatal(err)
}
if len(inputs) != 3 {
t.Fatalf("input count=%d want 3", len(inputs))
}
if inputs[0].Kind != "" || inputs[0].Value != "example.com" {
t.Fatalf("legacy rule changed: %+v", inputs[0])
}
if inputs[1].Kind != "icp" || inputs[2].Kind != "keyword" {
t.Fatalf("structured rules changed: %+v", inputs)
}
}
func TestCreateCompanyRejectsNormalizedDuplicateWithoutChangingScope(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("HTTP Strict Company %d", stamp)
companyID, _, _, _, _, err := m.pg.Companies().CreateCompanyWithScope(name, "", []db.ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("existing-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
s := &Server{m: m}
body, err := json.Marshal(map[string]any{
"name": " " + name + " ",
"scope": []map[string]string{{"kind": "domain", "value": fmt.Sprintf("new-%d.example", stamp)}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusConflict {
t.Fatalf("duplicate create status=%d body=%s", rec.Code, rec.Body.String())
}
scope, err := m.pg.Companies().GetScope(companyID)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("existing-%d.example", stamp) {
t.Fatalf("duplicate HTTP create changed existing scope: %+v", scope)
}
}
func TestDeleteCompanyRefreshesLiveTaskCompanyIDs(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete Company DTO %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("company deletion dto", "goal", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", bytes.NewBufferString(`{}`))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("delete company status=%d body=%s", rec.Code, rec.Body.String())
}
if dto := taskDTO(task, "created"); len(dto.CompanyIDs) != 0 {
t.Fatalf("live task DTO retained deleted company: %+v", dto.CompanyIDs)
}
persisted, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if persisted == nil || len(persisted.CompanyIDs) != 0 {
t.Fatalf("persisted task retained deleted company: %+v", persisted)
}
}
func TestCompanyScopeHTTPErrorClassificationAndBounds(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
t.Run("missing company is 404", func(t *testing.T) {
body := bytes.NewBufferString(`{"scope":[]}`)
req := httptest.NewRequest(http.MethodPost, "/api/companies/999999999/scope", body)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized Unicode scope is 400", func(t *testing.T) {
body, err := json.Marshal(map[string]any{
"name": "Oversized Scope",
"scope": []map[string]string{{
"kind": "keyword", "value": strings.Repeat("界", db.MaxCompanyScopeRawRunes+1),
}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("too many scope rules is 400", func(t *testing.T) {
scope := make([]map[string]string, db.MaxCompanyScopeRules+1)
for i := range scope {
scope[i] = map[string]string{"kind": "keyword", "value": fmt.Sprintf("keyword-%d", i)}
}
body, err := json.Marshal(map[string]any{"name": "Too Many Scope Rules", "scope": scope})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized request body is 413", func(t *testing.T) {
body := `{"name":"` + strings.Repeat("x", maxCompanyMutationBodyBytes+1) + `"}`
req := httptest.NewRequest(http.MethodPost, "/api/companies", strings.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
}
func TestListAssetsClassifiesValidationAndDatabaseErrors(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodGet, "/api/assets?dsl=(", nil)
rec := httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("invalid DSL status=%d body=%s", rec.Code, rec.Body.String())
}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req = httptest.NewRequest(http.MethodGet, "/api/assets?type=root_domain", nil)
rec = httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("database failure status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestDeleteCompanyRejectsBadJSONAndReportsMissing(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete JSON Company %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
for _, body := range []string{`{"delete_assets":`, `{}` + `{}`} {
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", strings.NewReader(body))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("body=%q status=%d response=%s", body, rec.Code, rec.Body.String())
}
company, err := m.pg.Companies().GetCompany(companyID)
if err != nil || company == nil {
t.Fatalf("bad JSON deleted company: company=%v err=%v", company, err)
}
}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/999999999", nil)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("missing delete status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestCompanyScopeSystemFailureIsHTTP500(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies/1/scope", bytes.NewBufferString(`{"scope":[]}`))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
}
server/conversations.go
+58 -12
View File
@@ -2,6 +2,7 @@ package server
import (
"context"
"errors"
"fmt"
"log"
"math"
@@ -10,12 +11,33 @@ import (
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/intercept"
)
const (
maxConversationRequestBytes = 64 << 10
maxConversationAgentKeyRunes = 120
maxConversationTitleRunes = 200
)
func decodeConversationRequest(w http.ResponseWriter, r *http.Request, value any) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxConversationRequestBytes)
if err := decode(r, value); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, err.Error())
}
return false
}
return true
}
// ---------- conversations (chat page) ----------
//
// A conversation is a ChatGPT-style thread bound to an agent key, independent of
@@ -46,8 +68,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id"`
}
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
if !decodeConversationRequest(w, r, &req) {
return
}
req.AgentKey = strings.TrimSpace(req.AgentKey)
@@ -55,6 +76,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "agent_key 不能为空")
return
}
if utf8.RuneCountInString(req.AgentKey) > maxConversationAgentKeyRunes {
writeErr(w, 400, fmt.Sprintf("agent_key 最多 %d 个字符", maxConversationAgentKeyRunes))
return
}
a, err := pg.GetAgentByKey(req.AgentKey)
if err != nil {
writeErr(w, 500, err.Error())
@@ -74,6 +99,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
if title == "" {
title = "新对话"
}
if utf8.RuneCountInString(title) > maxConversationTitleRunes {
writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes))
return
}
c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID)
if err != nil {
writeErr(w, 500, err.Error())
@@ -90,8 +119,7 @@ func (s *Server) pgUpdateConversation(w http.ResponseWriter, r *http.Request) {
var req struct {
LLMProfileID *int64 `json:"llm_profile_id"` // null clears the override
}
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
if !decodeConversationRequest(w, r, &req) {
return
}
if req.LLMProfileID != nil {
@@ -135,21 +163,39 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
if !ok {
return
}
var req struct{ Title string }
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
var req struct {
Title *string `json:"title"`
Pinned *bool `json:"pinned"`
}
if !decodeConversationRequest(w, r, &req) {
return
}
title := strings.TrimSpace(req.Title)
if title == "" {
writeErr(w, 400, "标题不能为空")
if req.Title == nil && req.Pinned == nil {
writeErr(w, 400, "至少需要提供 title 或 pinned")
return
}
if err := pg.RenameConversation(c.ID, title); err != nil {
if req.Title != nil {
title := strings.TrimSpace(*req.Title)
if title == "" {
writeErr(w, 400, "标题不能为空")
return
}
if utf8.RuneCountInString(title) > maxConversationTitleRunes {
writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes))
return
}
req.Title = &title
}
updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned})
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"ok": true})
if updated == nil {
writeErr(w, 404, "conversation not found")
return
}
writeJSON(w, 200, updated)
}
func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) {
server/dto.go
+9 -6
View File
@@ -51,6 +51,7 @@ type TaskDTO struct {
LLMFailoverState string `json:"llm_failover_state"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []string `json:"source_task_ids"`
CompanyIDs []int64 `json:"company_ids"`
}
// TokenTotalDTO is a whole-task (all agents) token aggregate.
@@ -71,9 +72,10 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
}
func taskDTO(t *Task, status string) TaskDTO {
lifecycle := t.lifecycleSnapshot()
llmState := t.llmStateSnapshot()
sourceIDs := make([]string, 0, len(t.SourceTaskIDs))
for _, id := range t.SourceTaskIDs {
sourceIDs := make([]string, 0, len(lifecycle.SourceTaskIDs))
for _, id := range lifecycle.SourceTaskIDs {
sourceIDs = append(sourceIDs, i64s(id))
}
profileIDs := append(make([]int64, 0, len(llmState.ProfileIDs)), llmState.ProfileIDs...)
@@ -85,16 +87,17 @@ func taskDTO(t *Task, status string) TaskDTO {
Status: status,
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
CreatedUnix: t.CreatedAt,
CompletedAt: completedRFC(t.CompletedAt),
CompletedUnix: t.CompletedAt,
Paused: t.Paused,
Queued: t.Queued,
CompletedAt: completedRFC(lifecycle.CompletedAt),
CompletedUnix: lifecycle.CompletedAt,
Paused: lifecycle.Paused,
Queued: lifecycle.Queued,
LLMProfileID: llmState.ProfileID,
LLMProfileIDs: profileIDs,
ActiveLLMProfileID: llmState.ActiveID,
LLMFailoverState: llmState.FailoverState,
LLMFailoverReason: llmState.FailoverReason,
SourceTaskIDs: sourceIDs,
CompanyIDs: lifecycle.CompanyIDs,
}
}
server/engine.go
+129 -40
View File
@@ -71,8 +71,11 @@ func preview(s string, n int) string {
const (
modelErrorRetries = 2 // model_error 收场后额外重试的次数
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
workControlWaitTimeout = 30 * time.Second
)
var errWorkControlConflict = errors.New("work control conflict")
// retryableWorkerModelError excludes errors already handled by the task router.
// In particular, a quota error after partial streaming advances the task cursor
// for the next LLM call but must not replay this whole intent on the backup.
@@ -101,10 +104,7 @@ type Engine struct {
llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
deleting sync.Map // taskID -> bool, delete barrier (no new task-owned writes)
// deleteWasPaused remembers whether BeginDelete temporarily paused an already
// running task. AbortDelete uses it to restore the exact pre-delete state.
deleteWasPaused sync.Map // taskID -> bool
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
// deleteMu makes installing the delete barrier atomic with registering a new
// task operation. Once BeginDelete returns, every admitted writer is reflected
@@ -161,7 +161,7 @@ type taskRuntime struct {
type workExecution struct {
cancel context.CancelCauseFunc
done chan struct{}
done chan error
action string // user action: pause | cancel
}
@@ -181,39 +181,33 @@ func (e *Engine) Pause(taskID string, cause error) {
}
// BeginDelete installs an execution barrier before task data/files are removed.
// The temporary pause is not a user pause: AbortDelete restores the exact state
// observed here when cleanup fails.
// The temporary pause is not a user pause. The server serializes this transition
// with lifecycle admission and tells AbortDelete whether the persisted task is
// paused/queued if cleanup fails.
func (e *Engine) BeginDelete(taskID string) bool {
e.deleteMu.Lock()
if _, loaded := e.deleting.LoadOrStore(taskID, true); loaded {
e.deleteMu.Unlock()
return false
}
wasPaused := e.IsPaused(taskID)
e.deleteWasPaused.Store(taskID, wasPaused)
e.paused.Store(taskID, true)
e.deleteMu.Unlock()
e.cancelExec(taskID, agent.AbortTaskDeleted)
return true
}
func (e *Engine) AbortDelete(taskID string) {
func (e *Engine) AbortDelete(taskID string, keepPaused bool) {
e.deleteMu.Lock()
if !e.IsDeleting(taskID) {
e.deleteMu.Unlock()
return
}
wasPaused := false
if v, ok := e.deleteWasPaused.Load(taskID); ok {
wasPaused, _ = v.(bool)
}
e.deleteWasPaused.Delete(taskID)
e.deleting.Delete(taskID)
if !wasPaused {
if !keepPaused {
e.paused.Delete(taskID)
}
e.deleteMu.Unlock()
if !wasPaused && e.m != nil {
if !keepPaused && e.m != nil {
if t, ok := e.m.Task(taskID); ok {
t.Notify()
}
@@ -291,8 +285,6 @@ func (e *Engine) StopTask(taskID string) {
e.stamped.Delete(taskID)
e.inflight.Delete(taskID)
e.coordStarted.Delete(taskID)
e.deleteWasPaused.Delete(taskID)
e.deleteMu.Lock()
e.deleting.Delete(taskID)
e.deleteMu.Unlock()
@@ -312,6 +304,17 @@ func (e *Engine) cancelExec(taskID string, cause error) {
// Resume un-pauses a task and nudges a fresh planning round. The next exec under
// it gets a fresh (uncancelled) context.
func (e *Engine) Resume(t *Task) {
// BeginDelete owns the pause barrier once deletion starts. A concurrent
// resume must never clear it and let a planner/worker re-enter while cleanup
// is waiting for task operations to drain.
if t == nil {
return
}
e.deleteMu.RLock()
defer e.deleteMu.RUnlock()
if e.IsDeleting(t.ID) {
return
}
e.paused.Delete(t.ID)
t.Notify()
}
@@ -393,14 +396,14 @@ func NewEngine(m *Manager) *Engine {
// registerWork records the cancel for the work currently running intentID.
func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) {
e.workMu.Lock()
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan struct{})}
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan error, 1)}
e.workMu.Unlock()
}
// detachWork removes the live control handle once Execute has returned. complete
// must be called after the final intent state write so a waiting cancel handler can
// safely delete the worker's blackboard output without racing a late write.
func (e *Engine) detachWork(intentID int64) (action string, complete func()) {
func (e *Engine) detachWork(intentID int64) (action string, complete func(error)) {
e.workMu.Lock()
run := e.work[intentID]
if run != nil {
@@ -413,27 +416,30 @@ func (e *Engine) detachWork(intentID int64) (action string, complete func()) {
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
e.steerMu.Unlock()
if run == nil {
return action, func() {}
return action, func(error) {}
}
return action, func() { close(run.done) }
return action, func(err error) { run.done <- err }
}
// ControlWork requests a user-visible pause or cancellation and waits until the
// worker has fully stopped writing. Cancellation cleanup is performed by the API
// handler after this returns; pause state is committed by runWorkerStep itself.
func (e *Engine) ControlWork(intentID int64, action string) error {
func (e *Engine) ControlWork(ctx context.Context, intentID int64, action string) error {
if action != "pause" && action != "cancel" {
return fmt.Errorf("unsupported work action %q", action)
}
if ctx == nil {
ctx = context.Background()
}
e.workMu.Lock()
run := e.work[intentID]
if run == nil {
e.workMu.Unlock()
return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID)
return fmt.Errorf("%w: 意图 %d 当前没有运行中的 work（可能已结束或未被领取）", errWorkControlConflict, intentID)
}
if run.action != "" && run.action != action {
if run.action != "" {
e.workMu.Unlock()
return fmt.Errorf("意图 %d 正在执行 %s 操作", intentID, run.action)
return fmt.Errorf("%w: 意图 %d 正在执行 %s 操作", errWorkControlConflict, intentID, run.action)
}
run.action = action
done := run.done
@@ -443,7 +449,41 @@ func (e *Engine) ControlWork(intentID int64, action string) error {
}
run.cancel(cause)
e.workMu.Unlock()
<-done
timer := time.NewTimer(workControlWaitTimeout)
defer timer.Stop()
select {
case err := <-done:
return err
case <-ctx.Done():
e.releaseWorkControl(intentID, run, action)
return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, ctx.Err())
case <-timer.C:
e.releaseWorkControl(intentID, run, action)
return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, context.DeadlineExceeded)
}
}
// releaseWorkControl drops only this caller's reservation after its wait is
// cancelled. The work context stays cancelled; runWorkerStep recognizes the
// named cancellation cause and settles the intent into the recoverable paused
// state even if the HTTP caller has gone away.
func (e *Engine) releaseWorkControl(intentID int64, run *workExecution, action string) {
e.workMu.Lock()
if current := e.work[intentID]; current == run && current.action == action {
current.action = ""
}
e.workMu.Unlock()
}
func transitionIntentState(store *db.ExplorationStore, intentID int64, expected, state string) error {
changed, err := store.CompareAndSetIntentState(intentID, expected, state)
if err != nil {
return err
}
if !changed {
return fmt.Errorf("%w: 意图 %d 不再是 %s 状态", db.ErrIntentStateConflict, intentID, expected)
}
return nil
}
@@ -744,7 +784,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
// terminal task (goals all met → done, or failed): the run is over. A
// resume/nudge — e.g. auto-resume of the active task on restart — must NOT
// re-plan (it would burn an LLM round and re-confirm a settled result).
if isTerminalStatus(t.Status) {
if isTerminalStatus(t.lifecycleSnapshot().Status) {
return
}
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
@@ -869,7 +909,9 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
emit := func(r db.Activity) { e.emitActivity(t, r) }
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
if ectx.Err() != nil || e.IsDeleting(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "open")
if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 领取后回退失败: %v", name, t.ID, intent.ID, err)
}
return true
}
// per-work child context so the planner's kill_work can stop just this work.
@@ -909,18 +951,41 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
// unregister would always be true (unregister cancels it) → every completed
// work would be wrongly marked stopped.
workCause := context.Cause(workCtx)
killed := workCtx.Err() != nil && ectx.Err() == nil
action, completeWork := e.detachWork(intent.ID)
defer completeWork()
// A caller may stop waiting and release its in-memory reservation before the
// agent honors cancellation. The named context cause remains authoritative and
// still settles the stopped run into a recoverable state.
if action == "" {
switch {
case errors.Is(workCause, agent.AbortWorkPausedByUser):
action = "pause"
case errors.Is(workCause, agent.AbortWorkCancelledByUser):
action = "cancel"
}
}
var controlErr error
defer func() { completeWork(controlErr) }()
if action == "pause" {
_ = t.Store.SetIntentState(intent.ID, "paused")
controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused")
if controlErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 暂停状态落库失败: %v", name, t.ID, intent.ID, controlErr)
return true
}
log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID)
e.touch(t.ID)
return true
}
if action == "cancel" {
// Leave the node in running until the waiting API handler deletes the intent
// and its direct outputs transactionally. No task Notify is emitted here.
// Park the stopped run in paused before handing cleanup to the API. If the
// request disconnects after cancellation, the intent remains recoverable and
// a later cancel can finish cleanup instead of leaving a phantom running row.
controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused")
if controlErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 取消栅栏落库失败: %v", name, t.ID, intent.ID, controlErr)
return true
}
log.Printf("[worker %s] task %s 意图 #%d 已停止，等待取消清理", name, t.ID, intent.ID)
e.touch(t.ID)
return true
@@ -928,14 +993,18 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// if a pause cancelled this run mid-flight, return the intent to the frontier
// so it is re-claimed on resume — the worker will resume the prior LLM
// conversation from its transcript instead of restarting from scratch.
if ectx.Err() != nil && e.IsPaused(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "open")
if ectx.Err() != nil && taskExecutionPaused(context.Cause(ectx)) {
if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 任务暂停回退失败: %v", name, t.ID, intent.ID, err)
}
return true
}
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
if ectx.Err() != nil && e.isSettling(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "exhausted")
if err := transitionIntentState(t.Store, intent.ID, "running", "exhausted"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 超时收尾状态落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted)，写回 %s", name, t.ID, intent.ID, wrote)
e.touch(t.ID)
return true
@@ -943,14 +1012,18 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
_ = t.Store.SetIntentState(intent.ID, "stopped")
if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 终态停止落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
e.touch(t.ID)
return true
}
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
if killed {
_ = t.Store.SetIntentState(intent.ID, "stopped")
if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d planner 停止落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
e.touch(t.ID)
t.Notify()
@@ -975,7 +1048,9 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
if state == "blocked" && isTaskLLMChainExhausted(err) {
_ = t.Store.SetIntentBlockedReason(intent.ID, db.IntentBlockedLLMQuota)
} else {
_ = t.Store.SetIntentState(intent.ID, state)
if stateErr := transitionIntentState(t.Store, intent.ID, "running", state); stateErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 终态 %s 落库失败: %v", name, t.ID, intent.ID, state, stateErr)
}
}
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote)
e.touch(t.ID)
@@ -983,6 +1058,20 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
return true
}
func taskExecutionPaused(cause error) bool {
var abort *agent.AbortCause
if !errors.As(cause, &abort) {
return false
}
switch abort.Code {
case "paused_by_user", "paused_by_orchestrator", "paused_on_reload", "paused_race_guard",
"queued_for_admission", "llm_unavailable_queued", "task_deleted":
return true
default:
return false
}
}
func sleepCtx(ctx context.Context, d time.Duration) (done bool) {
select {
case <-ctx.Done():
server/engine_cancelcause_test.go
+111 -3
View File
@@ -2,7 +2,9 @@ package server
import (
"context"
"errors"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
)
@@ -33,7 +35,7 @@ func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) {
t.Fatal("kill_work cancelled the entire task context")
}
_, complete := e.detachWork(42)
complete()
complete(nil)
}
func TestControlWorkCarriesUserActionCause(t *testing.T) {
@@ -49,7 +51,7 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
done := make(chan error, 1)
go func() { done <- e.ControlWork(42, tc.action) }()
go func() { done <- e.ControlWork(context.Background(), 42, tc.action) }()
<-ctx.Done()
if code, _, _, _ := agent.AbortReason(ctx); code != tc.code {
t.Fatalf("code=%q, want %q", code, tc.code)
@@ -58,7 +60,7 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
if action != tc.action {
t.Fatalf("action=%q, want %q", action, tc.action)
}
complete()
complete(nil)
if err := <-done; err != nil {
t.Fatal(err)
}
@@ -66,6 +68,61 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
}
}
func TestControlWorkReturnsSettlementError(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
done := make(chan error, 1)
go func() { done <- e.ControlWork(context.Background(), 42, "pause") }()
<-ctx.Done()
_, complete := e.detachWork(42)
want := errors.New("pause persistence failed")
complete(want)
if err := <-done; !errors.Is(err, want) {
t.Fatalf("ControlWork error=%v, want %v", err, want)
}
}
func TestControlWorkWaitHonorsContextAndReleasesReservation(t *testing.T) {
e := NewEngine(nil)
_, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
waitCtx, stop := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer stop()
if err := e.ControlWork(waitCtx, 42, "pause"); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("ControlWork error=%v, want deadline exceeded", err)
}
e.workMu.Lock()
run := e.work[42]
if run == nil || run.action != "" {
t.Fatalf("timed-out control left reservation: %+v", run)
}
e.workMu.Unlock()
_, complete := e.detachWork(42)
complete(nil)
}
func TestControlWorkRejectsConcurrentController(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
first := make(chan error, 1)
go func() { first <- e.ControlWork(context.Background(), 42, "pause") }()
<-ctx.Done()
if err := e.ControlWork(context.Background(), 42, "cancel"); !errors.Is(err, errWorkControlConflict) {
t.Fatalf("second control error=%v, want conflict", err)
}
action, complete := e.detachWork(42)
if action != "pause" {
t.Fatalf("winning action=%q, want pause", action)
}
complete(nil)
if err := <-first; err != nil {
t.Fatal(err)
}
}
func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) {
for _, tc := range []struct {
cause *agent.AbortCause
@@ -107,3 +164,54 @@ func TestResumeYieldsFreshContext(t *testing.T) {
t.Fatal("resume did not create a fresh execution context")
}
}
func TestTaskPauseCauseSurvivesImmediateResume(t *testing.T) {
e := NewEngine(nil)
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
old := e.execContextFor(context.Background(), task.ID)
e.Pause(task.ID, agent.AbortPausedByUser)
e.Resume(task)
if e.IsPaused(task.ID) {
t.Fatal("resume did not clear the task pause flag")
}
if old.Err() == nil || !taskExecutionPaused(context.Cause(old)) {
t.Fatalf("old execution cause=%v was not retained as a task pause", context.Cause(old))
}
if taskExecutionPaused(agent.AbortGoalMet) {
t.Fatal("goal completion was misclassified as a task pause")
}
}
func TestDeleteCancellationReturnsRunningIntentToFrontier(t *testing.T) {
if !taskExecutionPaused(agent.AbortTaskDeleted) {
t.Fatal("delete cancellation must settle a running intent as open so AbortDelete can resume it")
}
if taskExecutionPaused(agent.AbortGoalMet) {
t.Fatal("terminal task cancellation must not reopen completed work")
}
}
func TestResumeDoesNotClearDeleteBarrier(t *testing.T) {
e := NewEngine(nil)
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
e.paused.Store(task.ID, true)
e.deleting.Store(task.ID, true)
e.Resume(task)
if !e.IsPaused(task.ID) {
t.Fatal("resume cleared the pause barrier owned by task deletion")
}
}
func TestBootstrapAdmissionClearsQueuedPauseBarrier(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancel(context.Background())
cancel()
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
s := &Server{engine: e, ctx: ctx}
e.Pause(task.ID, agent.AbortPausedByUser)
s.startAdmittedTask(task, "bootstrap")
if e.IsPaused(task.ID) {
t.Fatal("bootstrap admission left the queued task behind the pause barrier")
}
}
server/engine_timeout.go
+16 -3
View File
@@ -71,13 +71,26 @@ func (e *Engine) taskDeadline(t *Task) int64 {
if v, ok := e.deadline.Load(t.ID); ok {
return v.(int64)
}
if t.DeadlineAt > 0 {
e.deadline.Store(t.ID, t.DeadlineAt)
return t.DeadlineAt
deadlineAt := t.lifecycleSnapshot().DeadlineAt
if deadlineAt > 0 {
e.deadline.Store(t.ID, deadlineAt)
return deadlineAt
}
return 0
}
// resetTimeoutRevival clears only the per-run timeout state after PostgreSQL has
// atomically committed timeout -> running and reset first_run_at/deadline_at. The
// configured TimeoutSeconds remains on Task, so the next real Planner/Worker run
// stamps a fresh full budget. coordStarted is reset because the coordinator that
// produced the timeout has already completed (or is in its final return path).
func (e *Engine) resetTimeoutRevival(taskID string) {
e.settling.Delete(taskID)
e.deadline.Delete(taskID)
e.stamped.Delete(taskID)
e.coordStarted.Delete(taskID)
}
// stampFirstRun records first_run_at + deadline_at on the FIRST real run (LLM ready)
// of a timeout task, once per process. No-op when the task has no timeout.
func (e *Engine) stampFirstRun(t *Task) {
server/findings_groups.go
+166
View File
@@ -0,0 +1,166 @@
package server
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"unicode/utf8"
"github.com/Autumn-27/artex/db"
)
const maxFindingFollowUpRunes = 4000
func findingPaginationParam(raw string, fallback, upperBound int) int {
value := atoiDefault(raw, fallback)
if value <= 0 {
value = fallback
}
if upperBound > 0 && value > upperBound {
value = upperBound
}
return value
}
// findingGroups serves the task-group page used by the global findings view.
// Findings inside a group continue to use the existing flat findings endpoint,
// preserving dashboard/export clients and giving every expanded group its own
// independent page cursor.
func (s *Server) findingGroups(w http.ResponseWriter, r *http.Request) {
q := r.URL.Query()
page := findingPaginationParam(q.Get("page"), 1, 0)
limit := findingPaginationParam(q.Get("limit"), 10, 100)
filter := db.FindingFilter{
Severity: normFilter(q.Get("severity")),
Status: normFilter(q.Get("status")),
VulnClass: normFilter(q.Get("vulnclass")),
TaskID: normFilter(q.Get("task_id")),
Sort: q.Get("sort"),
}
groups, total, findingTotal, err := s.m.pg.ListFindingGroups(filter, page, limit)
if err != nil {
writeErr(w, 500, err.Error())
return
}
// PostgreSQL owns paused/queued persistence. The live engine adds the same
// derived running/idle view used by task DTOs, so finding groups do not lag
// behind the task list while an Agent is actively executing.
if s.engine != nil {
for i := range groups {
if groups[i].TaskID == nil {
continue
}
if task, ok := s.m.Task(i64s(*groups[i].TaskID)); ok {
groups[i].TaskStatus = s.resolvedTaskStatus(task)
}
}
}
writeJSON(w, 200, map[string]any{
"items": groups,
"total": total,
"finding_total": findingTotal,
"page": page,
"page_size": limit,
})
}
// deepenFinding creates a high-priority human Worker intent under the finding's
// original task. The DB write validates and links the live finding node; task
// admission then revives or queues the task through the shared concurrency path.
func (s *Server) deepenFinding(w http.ResponseWriter, r *http.Request) {
id := int64(atoiDefault(r.PathValue("id"), 0))
if id <= 0 {
writeErr(w, 400, "bad finding id")
return
}
r.Body = http.MaxBytesReader(w, r.Body, 32<<10)
var req struct {
Description string `json:"description"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, "bad json: "+err.Error())
}
return
}
description := strings.TrimSpace(req.Description)
switch {
case description == "":
writeErr(w, 400, "description is required")
return
case utf8.RuneCountInString(description) > maxFindingFollowUpRunes:
writeErr(w, 400, fmt.Sprintf("description must be at most %d characters", maxFindingFollowUpRunes))
return
}
finding, err := s.m.pg.GetFinding(id)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if finding == nil {
writeErr(w, 404, "finding not found")
return
}
if finding.TaskID == nil || finding.NodeID == nil {
writeErr(w, 409, "finding origin task or node is no longer available")
return
}
t, ok := s.m.Task(i64s(*finding.TaskID))
if !ok || t == nil {
writeErr(w, 409, "finding origin task is no longer available")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "task is being deleted")
return
}
defer s.engine.decInflight(t.ID)
audit := db.Activity{
Worker: "system",
Kind: "text",
Summary: "人工提交漏洞深入利用意图",
Detail: description,
}
intentID, audit, err := t.Store.AddFindingFollowUpIntent(id, *finding.NodeID, description, audit)
if errors.Is(err, db.ErrFindingOriginUnavailable) {
writeErr(w, 409, "finding origin task or node is no longer available")
return
}
if err != nil {
writeErr(w, 500, err.Error())
return
}
queued, err := s.admitTask(t, "resume")
if err != nil {
if rollbackErr := t.Store.DiscardOpenIntent(intentID); rollbackErr != nil {
err = errors.Join(err, fmt.Errorf("discard follow-up intent %d: %w", intentID, rollbackErr))
}
writeErr(w, 500, err.Error())
return
}
// The audit row committed atomically with the intent. Publish that exact row;
// emitActivity would insert a second copy and break DB/SSE correspondence.
s.engine.Broadcaster().Publish(t.ID, audit)
s.engine.touch(t.ID)
// Admission can wake a Worker immediately. Report the persisted intent state
// observed at response time instead of always claiming it is still open. The
// task-level queued flag remains separate because a queued task intentionally
// keeps its Worker intent open until it receives a concurrency slot.
state := "open"
if current, stateErr := t.Store.GetNode(intentID); stateErr == nil && current != nil {
state = current.State
}
writeJSON(w, 200, map[string]any{
"task_id": t.ID,
"intent_id": i64s(intentID),
"state": state,
"queued": queued,
})
}
server/findings_groups_test.go
+321
View File
@@ -0,0 +1,321 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestFindingPaginationParam(t *testing.T) {
for _, test := range []struct {
raw, name string
fallback, upper, want int
}{
{name: "missing", raw: "", fallback: 10, upper: 100, want: 10},
{name: "invalid", raw: "not-a-number", fallback: 10, upper: 100, want: 10},
{name: "zero", raw: "0", fallback: 10, upper: 100, want: 10},
{name: "negative", raw: "-2", fallback: 1, upper: 0, want: 1},
{name: "bounded", raw: "500", fallback: 10, upper: 100, want: 100},
{name: "valid", raw: "25", fallback: 10, upper: 100, want: 25},
} {
t.Run(test.name, func(t *testing.T) {
if got := findingPaginationParam(test.raw, test.fallback, test.upper); got != test.want {
t.Fatalf("findingPaginationParam(%q,%d,%d)=%d, want %d",
test.raw, test.fallback, test.upper, got, test.want)
}
})
}
}
func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
marker := fmt.Sprintf("__server_finding_groups_%d__", time.Now().UnixNano())
createdTaskIDs := []int64{}
defer func() {
_, _ = m.pg.Exec(`DELETE FROM findings WHERE vulnclass=$1`, marker)
for _, taskID := range createdTaskIDs {
_ = m.pg.DeleteTask(taskID)
}
}()
taskA, err := m.CreateTask("finding group A", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
createdTaskIDs = append(createdTaskIDs, mustTaskID(t, taskA.ID))
taskB, err := m.CreateTask("finding group B", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
createdTaskIDs = append(createdTaskIDs, mustTaskID(t, taskB.ID))
for i, seed := range []struct {
taskID string
severity string
}{
{taskA.ID, db.SeverityCritical},
{taskA.ID, db.SeverityHigh},
{taskB.ID, db.SeverityLow},
} {
if _, err := m.pg.AddFinding(mustTaskID(t, seed.taskID), 0, marker, "", seed.severity,
fmt.Sprintf("finding %d", i), "", "test", nil); err != nil {
t.Fatalf("AddFinding[%d]: %v", i, err)
}
}
s := &Server{m: m}
request := httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&sort=severity&page=1&limit=1", nil)
recorder := httptest.NewRecorder()
s.findingGroups(recorder, request)
if recorder.Code != http.StatusOK {
t.Fatalf("groups status=%d body=%s", recorder.Code, recorder.Body.String())
}
var page struct {
Items []db.FindingGroup `json:"items"`
Total int `json:"total"`
FindingTotal int `json:"finding_total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
}
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
t.Fatal(err)
}
if len(page.Items) != 1 || page.Total != 2 || page.FindingTotal != 3 || page.Page != 1 || page.PageSize != 1 {
t.Fatalf("unexpected grouped page: %+v", page)
}
if page.Items[0].TaskID == nil || *page.Items[0].TaskID != mustTaskID(t, taskA.ID) ||
page.Items[0].Count != 2 || page.Items[0].Critical != 1 || page.Items[0].High != 1 {
t.Fatalf("unexpected first group: %+v", page.Items[0])
}
request = httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&page=-4&limit=0", nil)
recorder = httptest.NewRecorder()
s.findingGroups(recorder, request)
if recorder.Code != http.StatusOK {
t.Fatalf("normalized groups status=%d body=%s", recorder.Code, recorder.Body.String())
}
page = struct {
Items []db.FindingGroup `json:"items"`
Total int `json:"total"`
FindingTotal int `json:"finding_total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
}{}
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
t.Fatal(err)
}
if page.Page != 1 || page.PageSize != 10 || len(page.Items) != 2 {
t.Fatalf("pagination was not normalized: %+v", page)
}
}
func TestDeepenFindingCreatesAuditedIntentAndRevivesTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
task, err := m.CreateTask("deepen finding", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
store := task.Store
findingNodeID, err := store.AddNode(db.KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer m.pg.DeleteTask(taskID)
findingID, err := m.pg.AddFinding(taskID, findingNodeID, "test", "source", db.SeverityHigh, "source", "", "worker", nil)
if err != nil {
t.Fatal(err)
}
defer m.pg.DeleteFinding(findingID)
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel() // admission is exercised without starting background planner/worker work
s := &Server{m: m, engine: NewEngine(m), ctx: ctx}
live, unsubscribe := s.engine.Broadcaster().Subscribe(task.ID)
defer unsubscribe()
body := bytes.NewBufferString(`{"description":"验证完整利用链并保留可复现证据"}`)
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", body)
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("deepen status=%d body=%s", rec.Code, rec.Body.String())
}
var response struct {
TaskID string `json:"task_id"`
IntentID string `json:"intent_id"`
State string `json:"state"`
Queued bool `json:"queued"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &response); err != nil {
t.Fatal(err)
}
intentID, err := strconv.ParseInt(response.IntentID, 10, 64)
if err != nil || intentID <= 0 || response.TaskID != task.ID || response.State != "open" || response.Queued {
t.Fatalf("unexpected response: %+v", response)
}
intent, err := store.GetNode(intentID)
if err != nil || intent == nil {
t.Fatalf("intent=%+v err=%v", intent, err)
}
if intent.Priority != 10 || intent.Origin != "human" || intent.State != "open" {
t.Fatalf("intent metadata: %+v", intent)
}
if task.Status != "running" || task.Paused || task.Queued {
t.Fatalf("task was not revived through admission: %+v", task)
}
activity, _, err := store.ActivityList(&intentID, 0, 50)
if err != nil {
t.Fatal(err)
}
var persistedAudit db.Activity
auditCount := 0
for _, item := range activity {
if item.Worker == "system" && strings.Contains(item.Summary, "人工提交漏洞深入利用意图") {
persistedAudit = item
auditCount++
}
}
if auditCount != 1 || persistedAudit.ID <= 0 {
t.Fatalf("atomic system audit count=%d activity=%+v", auditCount, activity)
}
select {
case broadcast := <-live:
if broadcast.ID != persistedAudit.ID || broadcast.NodeID == nil || *broadcast.NodeID != intentID ||
broadcast.CreatedAt.IsZero() {
t.Fatalf("broadcast must reuse committed activity: broadcast=%+v persisted=%+v", broadcast, persistedAudit)
}
default:
t.Fatal("committed follow-up audit was not broadcast")
}
select {
case duplicate := <-live:
t.Fatalf("follow-up audit was broadcast more than once: %+v", duplicate)
default:
}
// The retained finding can no longer be deepened after its task and node are
// deleted; this must be a conflict rather than silently creating orphan work.
if err := m.pg.DeleteTask(taskID); err != nil {
t.Fatal(err)
}
req = httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", bytes.NewBufferString(`{"description":"retry"}`))
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec = httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusConflict {
t.Fatalf("deleted origin status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestDeepenFindingValidatesDescription(t *testing.T) {
s := &Server{}
tooLong, _ := json.Marshal(map[string]string{"description": strings.Repeat("x", maxFindingFollowUpRunes+1)})
for _, body := range []string{`{}`, `{"description":" "}`, string(tooLong)} {
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", bytes.NewBufferString(body))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("body=%s status=%d response=%s", body, rec.Code, rec.Body.String())
}
}
oversized := `{"description":"` + strings.Repeat("x", (32<<10)+1) + `"}`
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", strings.NewReader(oversized))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized status=%d response=%s", rec.Code, rec.Body.String())
}
}
func TestDeepenAdmissionFailureDiscardsFollowUpIntent(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("deepen rollback", "leave no orphan intent", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer func() {
_, _ = m.pg.Exec(`UPDATE tasks SET deleted_at=NULL WHERE id=$1`, taskID)
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
findingNodeID, err := task.Store.AddNode(db.KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
findingID, err := m.pg.AddFinding(taskID, findingNodeID, "test", "source", db.SeverityHigh, "source", "", "worker", nil)
if err != nil {
t.Fatal(err)
}
before, err := task.Store.ListByKind(db.KindIntent, 1000)
if err != nil {
t.Fatal(err)
}
var activitiesBefore int
if err := m.pg.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.Store.ID()).Scan(&activitiesBefore); err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
// Keep the in-memory handle and graph available while forcing the atomic
// admission UPDATE to reject this soft-deleted row.
if _, err := m.pg.Exec(`UPDATE tasks SET deleted_at=now() WHERE id=$1`, taskID); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
s := &Server{m: m, engine: NewEngine(m), ctx: ctx}
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen",
bytes.NewBufferString(`{"description":"this intent must be rolled back"}`))
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
after, err := task.Store.ListByKind(db.KindIntent, 1000)
if err != nil {
t.Fatal(err)
}
if len(after) != len(before) {
t.Fatalf("failed admission left an orphan follow-up: before=%d after=%d nodes=%+v", len(before), len(after), after)
}
var activitiesAfter int
if err := m.pg.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.Store.ID()).Scan(&activitiesAfter); err != nil {
t.Fatal(err)
}
if activitiesAfter != activitiesBefore {
t.Fatalf("failed admission left an orphan audit: before=%d after=%d", activitiesBefore, activitiesAfter)
}
}
server/goals.go
+239 -49
View File
@@ -33,14 +33,10 @@ func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
if seedFirstIntent {
s.seedFirstIntent(t)
}
if s.queueIfAtCapacity(t) {
s.engine.decInflight(t.ID)
return
s.engine.decInflight(t.ID)
if _, err := s.admitTask(t, "bootstrap"); err != nil {
log.Printf("[concurrency] task %s 启动失败: %v", t.ID, err)
}
go func() {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}()
}
func (s *Server) startTaskEngine(t *Task) {
@@ -70,7 +66,21 @@ func (s *Server) startTaskEngine(t *Task) {
}
func (s *Server) occupiesConcurrencySlot(t *Task) bool {
return t != nil && !t.Queued && !t.Paused && !s.engine.IsPaused(t.ID) && !isTerminalStatus(t.Status)
if t == nil {
return false
}
lifecycle := t.lifecycleSnapshot()
if lifecycle.Queued || lifecycle.Paused || isTerminalStatus(lifecycle.Status) {
return false
}
// Deletion temporarily pauses the Engine but has not committed yet. Preserve
// the task's slot until PostgreSQL deletion succeeds; otherwise FIFO promotion
// during the drain window could over-admit if deletion later aborts and the
// persisted running task is restored.
if s.engine.IsDeleting(t.ID) {
return true
}
return !s.engine.IsPaused(t.ID) && (s.engine.ReadyFor(t) || s.engine.ActiveLLMCalls(t.ID) > 0)
}
func (s *Server) runningTaskCount(excludeID string) int {
@@ -83,58 +93,255 @@ func (s *Server) runningTaskCount(excludeID string) int {
return count
}
func (s *Server) queueIfAtCapacity(t *Task) bool {
// admitTask is the single admission path for new, resumed, rerun and follow-up
// work. It atomically either starts the task or appends it to the persistent FIFO
// queue. mode is bootstrap for a freshly-created task and resume otherwise.
func (s *Server) admitTask(t *Task, mode string) (queued bool, err error) {
return s.admitTaskWhen(t, mode, false)
}
// admitPausedTask is the task-control resume path. The paused precondition is
// checked under the same scheduler lock as admission so a concurrent pause,
// dequeue or FIFO promotion cannot leave database and Engine state divergent.
func (s *Server) admitPausedTask(t *Task) (queued bool, err error) {
return s.admitTaskWhen(t, "resume", true)
}
func (s *Server) admitTaskWhen(t *Task, mode string, requirePaused bool) (queued bool, err error) {
if t == nil {
return false, fmt.Errorf("task not found")
}
if mode != "bootstrap" {
mode = "resume"
}
s.concMu.Lock()
defer s.concMu.Unlock()
// Delete installs its barrier under concMu as well. Re-resolve after acquiring
// the lock so a request that captured a task pointer before successful deletion
// cannot revive that stale handle after StopTask clears its Engine maps.
current, exists := s.m.Task(t.ID)
if !exists || current != t || s.engine.IsDeleting(t.ID) {
return false, fmt.Errorf("task is being deleted")
}
if !s.engine.beginTaskOperation(t.ID) {
return false, fmt.Errorf("task is being deleted")
}
defer s.engine.decInflight(t.ID)
lifecycle := t.lifecycleSnapshot()
if requirePaused {
if isTerminalStatus(lifecycle.Status) {
return false, fmt.Errorf("终态任务不能执行继续")
}
if !lifecycle.Paused {
return false, fmt.Errorf("仅已暂停的任务可以继续")
}
mode = s.resumeAdmissionMode(t)
}
wasTerminal := isTerminalStatus(lifecycle.Status)
wasPaused := lifecycle.Paused || s.engine.IsPaused(t.ID)
wasQueued := lifecycle.Queued
engineWasPaused := s.engine.IsPaused(t.ID)
enabled, limit := s.m.ConcurrencyLimit()
if !enabled || s.runningTaskCount(t.ID) < limit {
return false
ready := s.engine.ReadyFor(t)
// Work added to a task that is already admitted only needs a wake-up. This
// matters when the configured limit was lowered below the current running
// count: an existing task must not suddenly mark itself queued while its
// planner/workers are still live. Still compare-and-commit the persisted
// status: a concurrent terminal transition must win instead of being silently
// reported as a successful follow-up admission.
if !wasTerminal && !wasPaused && !wasQueued && s.engine.Started(t.ID) && (!enabled || ready) {
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, lifecycle.Status, false, "resume", false); err != nil {
return false, err
}
s.startAdmittedTask(t, "resume")
return false, nil
}
if err := s.m.SetTaskQueued(t.ID, true); err != nil {
log.Printf("[concurrency] task %s 置排队失败: %v", t.ID, err)
return false
// Preserve the original first-run mode across repeated admissions. Legacy
// queued rows have an empty queue_mode, so infer bootstrap from their graph.
if wasQueued {
switch lifecycle.QueueMode {
case "bootstrap":
mode = "bootstrap"
case "":
mode = s.resumeAdmissionMode(t)
}
}
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text",
Summary: fmt.Sprintf("已排队：达到并发上限 %d，等待空位后自动开始", limit)})
return true
readyBacklog := enabled && s.hasReadyQueuedTask(t.ID)
atCapacity := enabled && s.runningTaskCount(t.ID) >= limit
shouldQueue := enabled && (!ready || readyBacklog || atCapacity)
// Install the execution barrier before reviving a terminal/paused task. Without
// this ordering, its already-running worker loops can claim the newly-opened
// intent in the gap between status=running and queued=true.
if shouldQueue || wasTerminal || wasPaused || wasQueued {
s.engine.Pause(t.ID, agent.Causef("queued_for_admission", "任务等待运行准入",
"任务正在等待并发队列或准入状态提交，本次执行已停止；只有获得运行槽后才会重新领取意图"))
}
status := lifecycle.Status
if wasTerminal {
status = "running"
}
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, status, shouldQueue, mode, wasQueued); err != nil {
if !engineWasPaused && !wasPaused && !wasQueued {
s.engine.Resume(t)
}
return false, err
}
if lifecycle.Status == "timeout" && status == "running" {
// ApplyTaskAdmission reset this timed-out run's persisted clock. Clear the
// matching Engine gates before either parking it in FIFO or starting work;
// otherwise the old settling flag would make every worker skip forever.
s.engine.resetTimeoutRevival(t.ID)
}
if shouldQueue {
if !wasQueued {
summary := fmt.Sprintf("已排队：达到并发上限 %d，等待空位后自动开始", limit)
switch {
case !ready:
summary = "已排队：当前没有可运行的 LLM 配置，配置恢复后自动开始"
case readyBacklog:
summary = "已排队：已有更早的任务等待运行，将按 FIFO 顺序自动开始"
}
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: summary})
}
return true, nil
}
s.startAdmittedTask(t, mode)
return false, nil
}
func (s *Server) hasReadyQueuedTask(excludeID string) bool {
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if task.ID != excludeID && lifecycle.Queued && !lifecycle.Paused && !isTerminalStatus(lifecycle.Status) && s.engine.ReadyFor(task) {
return true
}
}
return false
}
func (s *Server) startAdmittedTask(t *Task, mode string) {
if mode == "bootstrap" {
if !s.engine.beginTaskOperation(t.ID) {
return
}
// A first-run task may have kept the Engine pause barrier while waiting
// in the concurrency queue. Clear it only after operation admission, or
// startTaskEngine's first execContextFor call would return a cancelled
// context and silently strand the dequeued task.
if s.engine.IsPaused(t.ID) {
s.engine.Resume(t)
}
go func() {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}()
return
}
s.engine.Run(s.ctx, t)
s.engine.Resume(t)
// Run returns early for an already-started task. Explicitly ensure a timeout
// coordinator exists after resume; resetTimeoutRevival cleared the completed
// coordinator's marker and the next real call will stamp a fresh deadline.
s.engine.startDeadlineCoordinator(s.ctx, t)
}
func (s *Server) reconcileConcurrency() {
s.concMu.Lock()
defer s.concMu.Unlock()
queued := []*Task{}
for _, task := range s.m.List() {
if task.Queued && !isTerminalStatus(task.Status) {
queued = append(queued, task)
enabled, limit := s.m.ConcurrencyLimit()
// A task whose provider chain becomes unavailable cannot do useful work and
// must not reserve a limited running slot forever. Park it persistently so a
// later profile edit/recovery can re-enter through the same FIFO path.
if enabled {
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if lifecycle.Queued || lifecycle.Paused || s.engine.IsPaused(task.ID) || s.engine.IsDeleting(task.ID) ||
isTerminalStatus(lifecycle.Status) || s.engine.ReadyFor(task) || s.engine.ActiveLLMCalls(task.ID) > 0 {
continue
}
mode := s.resumeAdmissionMode(task)
s.engine.Pause(task.ID, agent.Causef("llm_unavailable_queued", "LLM 不可用，任务进入等待队列",
"任务当前无法解析可运行的 Planner/Worker LLM，已释放并发槽；配置恢复后按队列顺序继续"))
if err := s.m.EnqueueTask(task.ID, mode); err != nil {
s.engine.Resume(task)
log.Printf("[concurrency] task %s 因 LLM 不可用入队失败: %v", task.ID, err)
continue
}
s.engine.emitActivity(task, db.Activity{Worker: "system", Kind: "text",
Summary: "已排队：当前没有可运行的 LLM 配置，配置恢复后自动开始"})
}
}
sort.Slice(queued, func(i, j int) bool { return queued[i].CreatedAt < queued[j].CreatedAt })
enabled, limit := s.m.ConcurrencyLimit()
type queuedTask struct {
task *Task
lifecycle taskLifecycleState
}
queued := []queuedTask{}
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
queued = append(queued, queuedTask{task: task, lifecycle: lifecycle})
}
}
sort.SliceStable(queued, func(i, j int) bool {
left, right := queued[i].lifecycle.QueuedAt, queued[j].lifecycle.QueuedAt
if left == 0 {
left = queued[i].task.CreatedAt * int64(1e9)
}
if right == 0 {
right = queued[j].task.CreatedAt * int64(1e9)
}
if left != right {
return left < right
}
// Old rows may not have queued_at and task creation timestamps are only
// kept to second precision in memory. Task ids are monotonic, providing a
// deterministic oldest-first fallback for those ties.
leftID, leftErr := strconv.ParseInt(queued[i].task.ID, 10, 64)
rightID, rightErr := strconv.ParseInt(queued[j].task.ID, 10, 64)
if leftErr == nil && rightErr == nil {
return leftID < rightID
}
return queued[i].task.ID < queued[j].task.ID
})
slots := len(queued)
if enabled {
slots = limit - s.runningTaskCount("")
}
for _, task := range queued {
for _, entry := range queued {
if slots <= 0 {
break
}
task := entry.task
lifecycle := task.lifecycleSnapshot()
// Keep FIFO order, but do not consume a concurrency slot for a task
// whose explicit chain/global provider is unavailable. It will be retried
// after the user configures or resets its LLM chain.
if task.Paused || !s.engine.ReadyFor(task) {
if lifecycle.Paused || (enabled && !s.engine.ReadyFor(task)) {
continue
}
mode := lifecycle.QueueMode
if mode != "bootstrap" && mode != "resume" {
mode = s.resumeAdmissionMode(task)
}
if mode != "bootstrap" {
mode = "resume"
}
if !s.engine.beginTaskOperation(task.ID) {
continue
}
if err := s.m.SetTaskQueued(task.ID, false); err != nil {
if err := s.m.ApplyTaskAdmission(task.ID, lifecycle.Status, lifecycle.Status, false, mode, false); err != nil {
s.engine.decInflight(task.ID)
continue
}
go func(t *Task) {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}(task)
s.startAdmittedTask(task, mode)
s.engine.decInflight(task.ID)
slots--
}
}
@@ -149,25 +356,8 @@ func (s *Server) reviveTask(t *Task) {
if t == nil {
return
}
if t.Queued {
return
}
// 终态 → 拉回 running(SetTaskStatus 会清 completed_at 并同步内存 handle 的 Status,
// 使 planner/worker 循环的终态门放行)。
if isTerminalStatus(t.Status) {
if err := s.m.SetTaskStatus(t.ID, "running"); err != nil {
log.Printf("[revive] task %s 置 running 失败: %v", t.ID, err)
}
}
// 确保引擎循环存活:已 started 只会 Notify;未 started(重启后未恢复的终态任务)则重起循环。
s.engine.Run(s.ctx, t)
// 暂停中 → 解除暂停(清引擎内存态 + Notify),并持久化 paused=false(存活过重启)。
if t.Paused || s.engine.IsPaused(t.ID) {
s.engine.Resume(t)
t.Paused = false
if err := s.m.SetTaskPaused(t.ID, false); err != nil {
log.Printf("[revive] task %s 解除暂停失败: %v", t.ID, err)
}
if _, err := s.admitTask(t, "resume"); err != nil {
log.Printf("[revive] task %s 恢复失败: %v", t.ID, err)
}
}
server/manager.go
+375 -52
View File
@@ -1,6 +1,7 @@
package server
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
@@ -26,21 +27,28 @@ import (
// sharing the process-wide asset store. ID is the PG task id as a string; ExpID
// is the exploration the task owns.
type Task struct {
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
Paused bool `json:"paused"`
Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
Paused bool `json:"paused"`
Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑
// QueuedAt is an internal Unix-nanosecond ordering key. It is deliberately
// finer than CreatedAt so several tasks enqueued in the same second retain
// their real FIFO order.
QueuedAt int64 `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
LLMChainRevision int64 `json:"-"`
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"`
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态；空/其它则由运行态推导)
// 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。
TimeoutSeconds int `json:"timeout_seconds"`
@@ -50,6 +58,7 @@ type Task struct {
Store *pgdb.ExplorationStore `json:"-"`
Guard *guard.Guard `json:"-"`
notify chan struct{}
lifecycleMu sync.RWMutex
llmMu sync.RWMutex
// pendingTriggers accumulates the concrete changes (worker done / finding) that
@@ -59,10 +68,72 @@ type Task struct {
pendingTriggers []agent.TriggerEvent
}
// taskLifecycleState is an internally consistent view of the mutable task
// lifecycle and inherited-scope context. Callers must use lifecycleSnapshot and
// updateLifecycle instead of reading or writing the corresponding Task fields
// directly after the task has been published by Manager.
type taskLifecycleState struct {
Status string
Paused bool
Queued bool
QueuedAt int64
QueueMode string
CompletedAt int64
FirstRunAt int64
DeadlineAt int64
SourceTaskIDs []int64
CompanyIDs []int64
}
func (t *Task) lifecycleSnapshot() taskLifecycleState {
if t == nil {
return taskLifecycleState{}
}
t.lifecycleMu.RLock()
defer t.lifecycleMu.RUnlock()
return t.lifecycleSnapshotLocked()
}
func (t *Task) lifecycleSnapshotLocked() taskLifecycleState {
return taskLifecycleState{
Status: t.Status,
Paused: t.Paused,
Queued: t.Queued,
QueuedAt: t.QueuedAt,
QueueMode: t.QueueMode,
CompletedAt: t.CompletedAt,
FirstRunAt: t.FirstRunAt,
DeadlineAt: t.DeadlineAt,
SourceTaskIDs: append([]int64(nil), t.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), t.CompanyIDs...),
}
}
func (t *Task) updateLifecycle(update func(*taskLifecycleState)) {
if t == nil || update == nil {
return
}
t.lifecycleMu.Lock()
state := t.lifecycleSnapshotLocked()
update(&state)
t.Status = state.Status
t.Paused = state.Paused
t.Queued = state.Queued
t.QueuedAt = state.QueuedAt
t.QueueMode = state.QueueMode
t.CompletedAt = state.CompletedAt
t.FirstRunAt = state.FirstRunAt
t.DeadlineAt = state.DeadlineAt
t.SourceTaskIDs = append(t.SourceTaskIDs[:0], state.SourceTaskIDs...)
t.CompanyIDs = append(t.CompanyIDs[:0], state.CompanyIDs...)
t.lifecycleMu.Unlock()
}
type taskLLMState struct {
ProfileID *int64
ProfileIDs []int64
ActiveID *int64
ChainRevision int64
FailoverState string
FailoverReason string
}
@@ -82,19 +153,25 @@ func (t *Task) llmStateSnapshot() taskLLMState {
ProfileID: cloneInt64Ptr(t.LLMProfileID),
ProfileIDs: append(make([]int64, 0, len(t.LLMProfileIDs)), t.LLMProfileIDs...),
ActiveID: cloneInt64Ptr(t.ActiveLLMProfileID),
ChainRevision: t.LLMChainRevision,
FailoverState: t.LLMFailoverState,
FailoverReason: t.LLMFailoverReason,
}
}
func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, state, reason string) {
func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, revision int64, state, reason string) bool {
t.llmMu.Lock()
defer t.llmMu.Unlock()
if revision < t.LLMChainRevision {
return false
}
t.LLMProfileID = cloneInt64Ptr(profileID)
t.ActiveLLMProfileID = cloneInt64Ptr(activeID)
t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...)
t.LLMChainRevision = revision
t.LLMFailoverState = state
t.LLMFailoverReason = reason
t.llmMu.Unlock()
return true
}
// DeleteTaskOptions controls cleanup of data stored outside the task's own
@@ -129,11 +206,16 @@ type Manager struct {
enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP)
interceptor *intercept.Interceptor // user-configured tool-call interception rules
mu sync.RWMutex
tasks map[string]*Task
active string
trafficOn bool // 流量捕获开关（默认关；settings.traffic_capture）
llmRecOn bool // LLM 录制开关（默认关；settings.llm_record）
companyMu sync.Mutex // serializes task/company-scope commits with live handle registration
// taskStateMu preserves commit order between PostgreSQL lifecycle writes and
// their in-memory mirrors. lifecycleMu makes snapshots race-free, but without
// this outer write lock an older request could commit first and publish last.
taskStateMu sync.Mutex
mu sync.RWMutex
tasks map[string]*Task
active string
trafficOn bool // 流量捕获开关（默认关；settings.traffic_capture）
llmRecOn bool // LLM 录制开关（默认关；settings.llm_record）
// 联网搜索开关与来源（默认关；settings.web_search_*）。brave-free 需要 braveKey；tavily 需要 tavilyKey。
// webSearchProxy 是独立出口代理(http/https/socks5)，与记录流量的 MITM 代理无关。
webSearchOn bool
@@ -594,15 +676,25 @@ func unixOrZero(t *time.Time) int64 {
return t.Unix()
}
func unixNanoOrZero(t *time.Time) int64 {
if t == nil {
return 0
}
return t.UnixNano()
}
func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task {
return &Task{
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued,
QueuedAt: unixNanoOrZero(pt.QueuedAt), QueueMode: pt.QueueMode,
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
LLMProfileID: pt.LLMProfileID,
LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID,
LLMChainRevision: pt.LLMChainRevision,
LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason,
SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), pt.CompanyIDs...),
TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds,
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
@@ -622,6 +714,10 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time
}
func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.TaskCreateOptions) (*Task, error) {
if len(opts.CompanyIDs) > 0 {
m.companyMu.Lock()
defer m.companyMu.Unlock()
}
pt, err := m.pg.CreateTaskWithOptions(description, goal, opts)
if err != nil {
return nil, err
@@ -634,6 +730,33 @@ func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.Task
return t, nil
}
// DeleteCompanyWithAssets keeps the database cascade and live task handles in
// one manager-level critical section. This closes the gap where a task could
// commit its company scope immediately before registration and miss the
// post-delete in-memory sweep.
func (m *Manager) DeleteCompanyWithAssets(id int64, deleteAssets bool) (int64, error) {
m.companyMu.Lock()
defer m.companyMu.Unlock()
assetsDeleted, err := m.pg.Companies().DeleteCompanyWithAssets(id, deleteAssets)
if err != nil {
return 0, err
}
m.mu.Lock()
for _, task := range m.tasks {
task.updateLifecycle(func(state *taskLifecycleState) {
companyIDs := make([]int64, 0, len(state.CompanyIDs))
for _, companyID := range state.CompanyIDs {
if companyID != id {
companyIDs = append(companyIDs, companyID)
}
}
state.CompanyIDs = companyIDs
})
}
m.mu.Unlock()
return assetsDeleted, nil
}
// ReplaceTaskLLMProfiles resets a non-terminal task's ordered provider chain and
// mirrors the committed state onto the live task handle.
func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activeProfileID int64) (int64, error) {
@@ -650,7 +773,7 @@ func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activePr
}
m.mu.Lock()
if task := m.tasks[id]; task != nil {
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason)
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMChainRevision, pt.LLMFailoverState, pt.LLMFailoverReason)
}
m.mu.Unlock()
if task, ok := m.Task(id); ok {
@@ -701,26 +824,212 @@ func (m *Manager) LoadExisting() []*Task {
// SetTaskPaused persists a task's paused state.
func (m *Manager) SetTaskPaused(id string, paused bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
return m.pg.SetPaused(n, paused)
}
// SetTaskQueued persists a task's queued (concurrency-hold) state and syncs the
// in-memory handle so listTasks/占位统计 see it immediately.
func (m *Manager) SetTaskQueued(id string, queued bool) error {
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if err := m.pg.SetQueued(n, queued); err != nil {
if err := m.pg.SetPaused(n, paused); err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Queued = queued
t.updateLifecycle(func(state *taskLifecycleState) {
state.Paused = paused
})
}
m.mu.Unlock()
return nil
}
// ApplyTaskAdmission atomically commits the lifecycle fields controlled by the
// concurrency scheduler. Keeping status, paused and queue metadata in one UPDATE
// prevents a failed resume from leaving a task half-revived (for example running
// but still user-paused, or dequeued without an Engine start).
//
// preservePosition applies only when the row is already queued. A repeated
// admission keeps its FIFO timestamp; a task that was explicitly paused and is
// now re-queued receives a fresh tail position.
func (m *Manager) ApplyTaskAdmission(id, expectedStatus, status string, queued bool, mode string, preservePosition bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if queued {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
} else {
mode = ""
}
var queuedAt, completedAt, firstRunAt, deadlineAt sql.NullTime
var committedMode string
err = m.pg.QueryRow(`UPDATE tasks
SET status=$2,
completed_at=CASE
WHEN $2 IN ('done','failed','timeout') THEN COALESCE(completed_at, now())
ELSE NULL
END,
paused=false,
queued=$3,
queued_at=CASE
WHEN NOT $3 THEN NULL
WHEN $5 AND queued THEN COALESCE(queued_at, now())
ELSE now()
END,
queue_mode=CASE
WHEN NOT $3 THEN ''
WHEN ($5 AND queued AND queue_mode='bootstrap') OR $4='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END,
first_run_at=CASE
WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL
ELSE first_run_at
END,
deadline_at=CASE
WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL
ELSE deadline_at
END
WHERE id=$1 AND deleted_at IS NULL AND status=$6
RETURNING queued_at, queue_mode, completed_at, first_run_at, deadline_at`, n, status, queued, mode, preservePosition, expectedStatus).
Scan(&queuedAt, &committedMode, &completedAt, &firstRunAt, &deadlineAt)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s lifecycle changed before admission (expected status %q)", id, expectedStatus)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
state.Paused = false
state.Queued = queued
state.QueueMode = committedMode
state.QueuedAt = 0
if queuedAt.Valid {
state.QueuedAt = queuedAt.Time.UnixNano()
}
state.CompletedAt = 0
if completedAt.Valid {
state.CompletedAt = completedAt.Time.Unix()
}
state.FirstRunAt = 0
if firstRunAt.Valid {
state.FirstRunAt = firstRunAt.Time.Unix()
}
state.DeadlineAt = 0
if deadlineAt.Valid {
state.DeadlineAt = deadlineAt.Time.Unix()
}
})
}
m.mu.Unlock()
return nil
}
// ApplyTaskPause atomically removes a task from the admission queue and records
// the user pause. queue_mode is intentionally retained so resuming a never-run
// bootstrap task still performs goal decomposition, but the next enqueue receives
// a new queued_at timestamp and therefore moves to the FIFO tail.
func (m *Manager) ApplyTaskPause(id string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
var mode string
err = m.pg.QueryRow(`UPDATE tasks
SET paused=true, queued=false, queued_at=NULL
WHERE id=$1 AND deleted_at IS NULL AND paused=false
AND status NOT IN ('done','failed','timeout')
RETURNING COALESCE(queue_mode,'')`, n).Scan(&mode)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s is unavailable for pause", id)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Paused = true
state.Queued = false
state.QueuedAt = 0
state.QueueMode = mode
})
}
m.mu.Unlock()
return nil
}
// EnqueueTask persists the concurrency hold and syncs the in-memory handle.
func (m *Manager) EnqueueTask(id, mode string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
var queuedAt time.Time
var committedMode string
err = m.pg.QueryRow(`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 AND deleted_at IS NULL
RETURNING queued_at, queue_mode`, n, mode).Scan(&queuedAt, &committedMode)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s is unavailable for enqueue", id)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.QueuedAt = queuedAt.UnixNano()
state.Queued = true
state.QueueMode = committedMode
})
}
m.mu.Unlock()
return nil
}
// DequeueTask removes the concurrency hold. clearMode=false is used when a user
// pauses a queued task so a later resume still knows whether bootstrap is needed.
func (m *Manager) DequeueTask(id string, clearMode bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if err := m.pg.Dequeue(n, clearMode); err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Queued = false
state.QueuedAt = 0
if clearMode {
state.QueueMode = ""
}
})
}
m.mu.Unlock()
return nil
@@ -731,7 +1040,7 @@ func (m *Manager) TaskStatus(id string) string {
m.mu.RLock()
defer m.mu.RUnlock()
if t := m.tasks[id]; t != nil {
return t.Status
return t.lifecycleSnapshot().Status
}
return ""
}
@@ -739,6 +1048,8 @@ func (m *Manager) TaskStatus(id string) string {
// StampTaskFirstRun stamps first_run_at + deadline_at on the first real run (idempotent
// in DB) and mirrors deadline_at on the live handle. Returns the deadline unix (0 = 不限).
func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return 0, err
@@ -759,10 +1070,12 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
if t.FirstRunAt == 0 {
t.FirstRunAt = time.Now().Unix()
}
t.DeadlineAt = dlUnix
t.updateLifecycle(func(state *taskLifecycleState) {
if state.FirstRunAt == 0 {
state.FirstRunAt = time.Now().Unix()
}
state.DeadlineAt = dlUnix
})
}
m.mu.Unlock()
return dlUnix, nil
@@ -772,6 +1085,8 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
// (resolves the completed↔timeout race — first terminal writer wins). Reflects the
// won status on the live handle. won=false means another terminal already stuck.
func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return false, err
@@ -782,10 +1097,12 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error)
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Status = status
if t.CompletedAt == 0 {
t.CompletedAt = time.Now().Unix()
}
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
if state.CompletedAt == 0 {
state.CompletedAt = time.Now().Unix()
}
})
}
m.mu.Unlock()
return true, nil
@@ -794,6 +1111,8 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error)
// SetTaskStatus persists a task's lifecycle status (e.g. "done") and reflects it
// on the in-memory handle so the derived DTO status shows it without a reload.
func (m *Manager) SetTaskStatus(id, status string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
@@ -803,16 +1122,18 @@ func (m *Manager) SetTaskStatus(id, status string) error {
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Status = status
// mirror the DB's completed_at stamp on the live handle so the DTO shows the
// finish time without a reload (terminal → stamp once; else clear).
if pgdb.IsTerminal(status) {
if t.CompletedAt == 0 {
t.CompletedAt = time.Now().Unix()
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
// Mirror the DB's completed_at stamp on the live handle so the DTO shows
// the finish time without a reload (terminal -> stamp once; else clear).
if pgdb.IsTerminal(status) {
if state.CompletedAt == 0 {
state.CompletedAt = time.Now().Unix()
}
} else {
state.CompletedAt = 0
}
} else {
t.CompletedAt = 0
}
})
}
m.mu.Unlock()
return nil
@@ -927,13 +1248,15 @@ func (m *Manager) forgetTask(id string, numericID int64) {
m.mu.Lock()
delete(m.tasks, id)
for _, task := range m.tasks {
kept := task.SourceTaskIDs[:0]
for _, sourceID := range task.SourceTaskIDs {
if sourceID != numericID {
kept = append(kept, sourceID)
task.updateLifecycle(func(state *taskLifecycleState) {
kept := make([]int64, 0, len(state.SourceTaskIDs))
for _, sourceID := range state.SourceTaskIDs {
if sourceID != numericID {
kept = append(kept, sourceID)
}
}
}
task.SourceTaskIDs = kept
state.SourceTaskIDs = kept
})
}
if m.active == id {
m.active = ""
server/manager_lifecycle_test.go
+109
View File
@@ -0,0 +1,109 @@
package server
import (
"fmt"
"sync"
"testing"
)
func TestTaskLifecycleSnapshotCopiesContextSlices(t *testing.T) {
task := &Task{
Status: "running",
SourceTaskIDs: []int64{11, 12},
CompanyIDs: []int64{21, 22},
}
snapshot := task.lifecycleSnapshot()
snapshot.SourceTaskIDs[0] = 99
snapshot.CompanyIDs[0] = 98
fresh := task.lifecycleSnapshot()
if got := fresh.SourceTaskIDs[0]; got != 11 {
t.Fatalf("snapshot source IDs alias task storage: got %d", got)
}
if got := fresh.CompanyIDs[0]; got != 21 {
t.Fatalf("snapshot company IDs alias task storage: got %d", got)
}
sources := []int64{31, 32}
companies := []int64{41, 42}
task.updateLifecycle(func(state *taskLifecycleState) {
state.SourceTaskIDs = sources
state.CompanyIDs = companies
})
sources[0] = 97
companies[0] = 96
fresh = task.lifecycleSnapshot()
if got := fresh.SourceTaskIDs[0]; got != 31 {
t.Fatalf("lifecycle update retained caller source slice: got %d", got)
}
if got := fresh.CompanyIDs[0]; got != 41 {
t.Fatalf("lifecycle update retained caller company slice: got %d", got)
}
}
func TestTaskLifecycleSnapshotConcurrentConsistency(t *testing.T) {
task := &Task{}
writeState := func(generation int64, queued bool) {
task.updateLifecycle(func(state *taskLifecycleState) {
state.Status = map[bool]string{true: "queued", false: "paused"}[queued]
state.Queued = queued
state.Paused = !queued
state.QueuedAt = generation
state.QueueMode = map[bool]string{true: "resume", false: "bootstrap"}[queued]
state.CompletedAt = generation
state.FirstRunAt = generation
state.DeadlineAt = generation
state.SourceTaskIDs = []int64{generation, generation + 1}
state.CompanyIDs = []int64{generation, generation + 1}
})
}
writeState(1, true)
var wg sync.WaitGroup
for writer := 0; writer < 4; writer++ {
writer := writer
wg.Add(1)
go func() {
defer wg.Done()
for i := 1; i <= 1000; i++ {
generation := int64(writer*1000 + i)
writeState(generation, generation%2 == 0)
}
}()
}
errCh := make(chan error, 4)
for reader := 0; reader < 4; reader++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 1000; i++ {
state := task.lifecycleSnapshot()
if state.Queued == state.Paused {
errCh <- fmt.Errorf("queued=%v paused=%v", state.Queued, state.Paused)
return
}
wantStatus := map[bool]string{true: "queued", false: "paused"}[state.Queued]
wantMode := map[bool]string{true: "resume", false: "bootstrap"}[state.Queued]
if state.Status != wantStatus || state.QueueMode != wantMode {
errCh <- fmt.Errorf("inconsistent status/mode: %+v", state)
return
}
generation := state.QueuedAt
if state.CompletedAt != generation || state.FirstRunAt != generation || state.DeadlineAt != generation ||
len(state.SourceTaskIDs) != 2 || len(state.CompanyIDs) != 2 ||
state.SourceTaskIDs[0] != generation || state.CompanyIDs[0] != generation {
errCh <- fmt.Errorf("mixed lifecycle snapshot: %+v", state)
return
}
}
}()
}
wg.Wait()
close(errCh)
for err := range errCh {
t.Fatal(err)
}
}
server/orchestration.go
+9 -6
View File
@@ -302,11 +302,13 @@ func (s *Server) toolPauseTask() actool.CoreTool {
TaskID string `json:"task_id"`
}
_ = json.Unmarshal(in, &a)
if _, ok := s.m.Task(a.TaskID); !ok {
t, ok := s.m.Task(a.TaskID)
if !ok {
return actool.Errorf("task 不存在: " + a.TaskID), nil
}
s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator)
s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask)
if _, err := s.applyTaskControlWithCause(t, "pause", agent.AbortPausedByOrchestrator); err != nil {
return actool.Errorf(err.Error()), nil
}
return actool.Text("task paused: " + a.TaskID), nil
})
}
@@ -416,10 +418,11 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
func (s *Server) deriveTaskStatus(t *Task) string {
lifecycle := t.lifecycleSnapshot()
switch {
case isTerminalStatus(t.Status):
return t.Status
case t.Paused || s.engine.IsPaused(t.ID):
case isTerminalStatus(lifecycle.Status):
return lifecycle.Status
case lifecycle.Paused || s.engine.IsPaused(t.ID):
return "paused"
case s.engine.ReadyFor(t) && s.engine.Started(t.ID):
return "running"
server/server.go
+152 -100
View File
@@ -4,7 +4,7 @@ import (
"archive/zip"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
@@ -238,26 +238,32 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
// reload tasks persisted on disk so the task list survives a restart, and
// restore persisted paused state (so a task paused before restart stays paused).
for _, t := range m.LoadExisting() {
lifecycle := t.lifecycleSnapshot()
// clear stale 'running' intents from a prior crash/restart (no live worker
// owns them) so they re-claim instead of spinning forever in the UI.
if n, _ := t.Store.ResetRunningIntents(); n > 0 {
log.Printf("[engine] task %s 重置 %d 个残留 running 意图为 open", t.ID, n)
}
if t.Paused {
if lifecycle.Paused {
s.engine.Pause(t.ID, agent.AbortPausedOnReload)
}
// 任务级超时:为每个未终态、带 timeout 的任务起 deadline 协调器,独立于 planner/worker
// loop——非活跃任务重启后也能在到点后被收尾(deadline 已过则立即走收尾时序)。
if !isTerminalStatus(t.Status) {
if !isTerminalStatus(lifecycle.Status) {
s.engine.startDeadlineCoordinator(ctx, t)
}
}
// resume the active task's engine; other tasks resume when opened (setActive).
// (a restored-paused task's loops still start but idle until resumed.)
// 排队中的任务不在此启动——留给 reconcileConcurrency 按空位补位。
if t := m.ActiveTask(); t != nil && !t.Queued {
s.engine.Run(ctx, t)
// Restore every task that had already been admitted before shutdown. Starting
// only the active UI task left other non-queued tasks counted as concurrency
// occupants without live loops, which could permanently block the persistent
// FIFO. Paused loops remain idle; queued tasks are admitted below as slots allow.
for _, t := range m.List() {
lifecycle := t.lifecycleSnapshot()
if !lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
s.engine.Run(ctx, t)
}
}
go s.reconcileConcurrency()
return s
}
@@ -636,7 +642,12 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/tasks", s.listTasks)
mux.HandleFunc("POST /api/tasks", s.createTask)
mux.HandleFunc("GET /api/task-templates", s.pgListTaskTemplates)
mux.HandleFunc("POST /api/task-templates", s.pgCreateTaskTemplate)
mux.HandleFunc("PATCH /api/task-templates/{id}", s.pgUpdateTaskTemplate)
mux.HandleFunc("DELETE /api/task-templates/{id}", s.pgDeleteTaskTemplate)
mux.HandleFunc("GET /api/tasks/{id}", s.getTask)
mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch)
mux.HandleFunc("GET /api/tasks/{id}/coverage", s.taskCoverage)
mux.HandleFunc("GET /api/tasks/{id}/coverage-graph", s.taskCoverageGraph)
mux.HandleFunc("GET /api/tasks/{id}/asset-refs", s.taskAssetRefs)
@@ -654,7 +665,9 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("DELETE /api/tasks/{id}/scope/{sid}", s.taskScopeDelete)
mux.HandleFunc("POST /api/tasks/{id}/control", s.control)
mux.HandleFunc("PUT /api/tasks/{id}/llm", s.updateTaskLLMProfiles)
mux.HandleFunc("GET /api/tasks/{id}/llm/resolution", s.taskLLMResolutionHandler)
mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/control", s.controlIntent)
mux.HandleFunc("POST /api/tasks/{id}/intents/control/batch", s.controlIntentsBatch)
mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/rerun", s.rerunIntent) // 重跑单条 blocked/exhausted/stopped 意图
mux.HandleFunc("POST /api/tasks/{id}/intents/rerun-blocked", s.rerunBlocked) // 批量重跑本任务全部 blocked 意图
mux.HandleFunc("POST /api/active", s.setActive)
@@ -679,10 +692,12 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/exploration/frontier", s.frontier)
mux.HandleFunc("GET /api/exploration/findings", s.findings)
mux.HandleFunc("GET /api/exploration/findings/groups", s.findingGroups)
mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats)
mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport)
mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding)
mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage)
mux.HandleFunc("POST /api/exploration/findings/{id}/deepen", s.deepenFinding)
mux.HandleFunc("PATCH /api/exploration/findings/{id}", s.patchFinding)
mux.HandleFunc("DELETE /api/exploration/findings/{id}", s.deleteFinding)
mux.HandleFunc("GET /api/exploration/intents", s.intents)
@@ -715,6 +730,7 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/report", s.getReport)
mux.HandleFunc("POST /api/chat", s.chat)
mux.HandleFunc("POST /api/chat/upload", s.chatUpload) // 方式1 文件上传:落到会话/任务工作目录 uploads/
mux.HandleFunc("GET /api/tasks/{id}/chat/status", s.taskChatStatus)
mux.HandleFunc("POST /api/tasks/{id}/chat/stop", s.stopChat)
// --- 管理后台 API (PostgreSQL 数据源; 新版数据库与管理后台方案) ---
@@ -948,12 +964,13 @@ func (s *Server) resolvedTaskStatus(t *Task) string {
if t == nil {
return "created"
}
lifecycle := t.lifecycleSnapshot()
switch {
case isTerminalStatus(t.Status):
return t.Status
case t.Queued:
case isTerminalStatus(lifecycle.Status):
return lifecycle.Status
case lifecycle.Queued:
return "queued"
case t.Paused || s.engine.IsPaused(t.ID):
case lifecycle.Paused || s.engine.IsPaused(t.ID):
return "paused"
case t.llmStateSnapshot().FailoverState == "chain_exhausted" && s.engine.Started(t.ID):
// LLM readiness is an execution dependency, not lifecycle state. Keeping
@@ -980,7 +997,7 @@ func (s *Server) setActive(w http.ResponseWriter, r *http.Request) {
}
// resume the engine for the opened task (idempotent — no-op if already running).
// 排队中的任务:仅设为活跃可查看,不启动引擎(维持并发上限,由 reconcile 补位)。
if t, ok := s.m.Task(req.ID); ok && !t.Queued {
if t, ok := s.m.Task(req.ID); ok && !t.lifecycleSnapshot().Queued {
s.engine.Run(s.ctx, t)
}
writeJSON(w, 200, map[string]any{"active": req.ID})
@@ -1000,28 +1017,16 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, err.Error())
return
}
switch req.Action {
case "pause":
s.engine.Pause(t.ID, agent.AbortPausedByUser)
// Main Agent turns run outside Engine's planner/worker context so they need
// their own cancellation: pausing terminates the CURRENTLY running turn. New
// turns are still allowed while paused — the main-agent console is independent
// of task pause (the chat handler no longer rejects on IsPaused).
s.cancelTaskChat(t.ID, agent.AbortChatPausedWithTask)
case "resume":
s.engine.Run(s.ctx, t) // ensure loops are alive, then un-pause + nudge
s.engine.Resume(t)
default:
if req.Action != "pause" && req.Action != "resume" {
writeErr(w, 400, "action must be pause|resume")
return
}
paused := s.engine.IsPaused(t.ID)
t.Paused = paused // keep the in-memory task (shown in /api/tasks) in sync
if err := s.m.SetTaskPaused(t.ID, paused); err != nil { // persist (survives restart)
log.Printf("[control] persist paused %s: %v", t.ID, err)
result, err := s.applyTaskControl(t, req.Action)
if err != nil {
writeErr(w, 409, err.Error())
return
}
log.Printf("[task] #%s %s", t.ID, map[bool]string{true: "已暂停", false: "已恢复"}[paused])
writeJSON(w, 200, map[string]any{"id": t.ID, "paused": paused})
writeJSON(w, 200, result)
}
// controlIntent pauses, resumes, or cancels one local worker intent. A running
@@ -1056,70 +1061,27 @@ func (s *Server) controlIntent(w http.ResponseWriter, r *http.Request) {
return
}
node, err := t.Store.GetNode(iid)
result, err := s.applyIntentControl(r.Context(), t, iid, req.Action)
if err != nil {
writeErr(w, 500, err.Error())
writeErr(w, 409, err.Error())
return
}
if node == nil {
if inherited, sourceErr := t.Store.GetNodeWithSources(iid); sourceErr == nil && inherited != nil && inherited.Inherited {
writeErr(w, 409, "继承意图为只读，不能控制")
} else {
writeErr(w, 404, "intent not found")
}
return
}
if node.Kind != db.KindIntent {
writeErr(w, 409, "node is not an intent")
return
}
switch req.Action {
case "pause":
if node.State != "running" {
writeErr(w, 409, "仅运行中的意图可以暂停")
return
}
if err := s.engine.ControlWork(iid, "pause"); err != nil {
writeErr(w, 409, err.Error())
return
}
writeJSON(w, 200, map[string]any{"id": iid, "state": "paused"})
case "resume":
if node.State != "paused" {
writeErr(w, 409, "仅已暂停的意图可以恢复")
return
}
if err := t.Store.SetIntentState(iid, "open"); err != nil {
writeErr(w, 500, err.Error())
return
}
t.Notify()
writeJSON(w, 200, map[string]any{"id": iid, "state": "open"})
case "cancel":
if node.State != "running" && node.State != "paused" {
writeErr(w, 409, "仅运行中或已暂停的意图可以取消")
return
}
if node.State == "running" {
if err := s.engine.ControlWork(iid, "cancel"); err != nil {
writeErr(w, 409, err.Error())
return
}
}
deleted, err := t.Store.CancelIntent(iid)
if err != nil {
writeErr(w, 409, err.Error())
return
}
t.Notify()
writeJSON(w, 200, map[string]any{"id": iid, "state": "cancelled", "deleted": deleted})
}
writeJSON(w, 200, result)
}
// rerunIntent 重跑一条没跑成功的意图(blocked/exhausted/stopped):把它置回 open,worker
// 会重新认领、从头再跑(已写回图谱的 fact/finding/asset 保留);若任务已终态/暂停则顺带复活。
// 用于「出错的 work 点击继续运行」——网络/LLM 抖动导致 blocked 后可一键重试。
func restoreRerunIntent(t *Task, before *db.Node) error {
if t == nil || before == nil {
return fmt.Errorf("missing intent rollback snapshot")
}
if before.State == "blocked" && before.BlockedReason != "" {
return t.Store.SetIntentBlockedReason(before.ID, before.BlockedReason)
}
return t.Store.SetIntentState(before.ID, before.State)
}
func (s *Server) rerunIntent(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
@@ -1131,6 +1093,16 @@ func (s *Server) rerunIntent(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "bad intent id")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除，无法重跑意图")
return
}
defer s.engine.decInflight(t.ID)
before, err := t.Store.GetNode(iid)
if err != nil {
writeErr(w, 500, err.Error())
return
}
reopened, err := t.Store.ReopenIntent(iid)
if err != nil {
writeErr(w, 500, err.Error())
@@ -1140,9 +1112,16 @@ func (s *Server) rerunIntent(w http.ResponseWriter, r *http.Request) {
writeErr(w, 409, "该意图不是可重跑状态(仅 blocked/exhausted/stopped 可重跑)")
return
}
s.reviveTask(t) // 终态/暂停 → 拉回 running,确保 worker 循环存活并重新认领
queued, err := s.admitTask(t, "resume")
if err != nil {
if rollbackErr := restoreRerunIntent(t, before); rollbackErr != nil {
err = fmt.Errorf("%w; restore intent %d after admission failure: %v", err, iid, rollbackErr)
}
writeErr(w, 500, err.Error())
return
}
log.Printf("[task] #%s 意图 #%d 已重开(重跑)", t.ID, iid)
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": iid})
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": iid, "queued": queued})
}
// rerunBlocked 批量重跑本任务全部 blocked 意图(适合一次网络/LLM 断连导致多条 blocked 后
@@ -1153,16 +1132,47 @@ func (s *Server) rerunBlocked(w http.ResponseWriter, r *http.Request) {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除，无法重跑意图")
return
}
defer s.engine.decInflight(t.ID)
intents, err := t.Store.ListByKind(db.KindIntent, 1000000)
if err != nil {
writeErr(w, 500, err.Error())
return
}
before := make([]*db.Node, 0)
for _, intent := range intents {
if intent.State == "blocked" {
copy := *intent
before = append(before, &copy)
}
}
n, err := t.Store.ReopenBlockedIntents()
if err != nil {
writeErr(w, 500, err.Error())
return
}
queued := false
if n > 0 {
s.reviveTask(t)
queued, err = s.admitTask(t, "resume")
if err != nil {
rollbackErrors := make([]string, 0)
for _, intent := range before {
if rollbackErr := restoreRerunIntent(t, intent); rollbackErr != nil {
rollbackErrors = append(rollbackErrors, fmt.Sprintf("intent %d: %v", intent.ID, rollbackErr))
}
}
if len(rollbackErrors) > 0 {
err = fmt.Errorf("%w; restore blocked intents after admission failure: %s", err, strings.Join(rollbackErrors, "; "))
}
writeErr(w, 500, err.Error())
return
}
log.Printf("[task] #%s 批量重开 %d 条 blocked 意图", t.ID, n)
}
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": n})
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": n, "queued": queued})
}
// getLLM returns the current LLM config (key never exposed).
@@ -1294,6 +1304,7 @@ type createTaskReq struct {
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效
SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务
CompanyIDs []int64 `json:"company_ids,omitempty"` // 关联企业资产范围;不复制资产或强制生成意图
TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600
SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。
@@ -1337,11 +1348,21 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) {
seenSources[id] = true
sourceIDs = append(sourceIDs, id)
}
companyIDs, err := db.NormalizeTaskCompanyIDs(req.CompanyIDs)
if err != nil {
writeErr(w, 400, fmt.Sprintf("关联企业无效：最多选择 %d 个有效企业", db.MaxTaskCompanyCount))
return
}
req.CompanyIDs = companyIDs
t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{
SourceTaskIDs: sourceIDs, LLMProfileIDs: req.LLMProfileIDs,
SourceTaskIDs: sourceIDs, CompanyIDs: req.CompanyIDs, LLMProfileIDs: req.LLMProfileIDs,
TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds,
})
if err != nil {
if errors.Is(err, db.ErrTaskCompanyIDsInvalid) || errors.Is(err, db.ErrTaskCompanyNotFound) {
writeErr(w, 400, "关联企业不存在或无效")
return
}
writeErr(w, 500, err.Error())
return
}
@@ -1372,7 +1393,7 @@ func (s *Server) updateTaskLLMProfiles(w http.ResponseWriter, r *http.Request) {
writeErr(w, 404, "task not found")
return
}
if isTerminalStatus(t.Status) {
if isTerminalStatus(t.lifecycleSnapshot().Status) {
writeErr(w, 409, "已结束的任务不能修改 LLM 配置")
return
}
@@ -1403,8 +1424,11 @@ func (s *Server) updateTaskLLMProfiles(w http.ResponseWriter, r *http.Request) {
return
}
// Profile edits affect only subsequent LLM calls. Existing in-flight calls
// retain their provider; a nudge lets an idle running task resume promptly.
// retain their provider. Concurrency reconciliation treats ActiveLLMCalls as a
// live slot and postpones an unavailable task's pause/queue transition until
// the call returns; a nudge lets an idle running task resume promptly.
t.Notify()
go s.reconcileConcurrency()
llmState := t.llmStateSnapshot()
result := map[string]any{
"id": t.ID, "llm_profile_ids": llmState.ProfileIDs,
@@ -1737,8 +1761,8 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) {
writeJSON(w, 200, out)
return
}
page := atoiDefault(q.Get("page"), 1)
limit := min(atoiDefault(q.Get("limit"), 20), 200)
page := findingPaginationParam(q.Get("page"), 1, 0)
limit := findingPaginationParam(q.Get("limit"), 20, 200)
filter := db.FindingFilter{
Severity: normFilter(q.Get("severity")),
Status: normFilter(q.Get("status")),
@@ -2023,7 +2047,7 @@ func findingProvenanceInTask(contextTask *Task, findingTaskID *int64) (sourceTas
if *findingTaskID == contextID {
return "", false, true
}
for _, sourceID := range contextTask.SourceTaskIDs {
for _, sourceID := range contextTask.lifecycleSnapshot().SourceTaskIDs {
if sourceID == *findingTaskID {
return i64s(sourceID), true, true
}
@@ -2498,7 +2522,11 @@ func (s *Server) activityHistory(w http.ResponseWriter, r *http.Request) {
func (s *Server) tokenStats(w http.ResponseWriter, r *http.Request) {
t := s.m.ResolveTask(r.URL.Query().Get("task"))
if t == nil {
writeJSON(w, 200, map[string]any{"workers": []any{}})
writeJSON(w, 200, map[string]any{
"workers": []db.TokenUsage{},
"sessions": []db.SessionTokenUsage{},
"total": tokenTotalDTO(db.TokenUsage{}),
})
return
}
stats, err := t.Store.TokenStatsByWorker()
@@ -2506,8 +2534,17 @@ func (s *Server) tokenStats(w http.ResponseWriter, r *http.Request) {
writeErr(w, 500, err.Error())
return
}
total, _ := t.Store.TokenTotal() // whole-task total (all agents)
writeJSON(w, 200, map[string]any{"workers": stats, "total": tokenTotalDTO(total)})
sessions, err := t.Store.TokenStatsBySession()
if err != nil {
writeErr(w, 500, err.Error())
return
}
total, err := t.Store.TokenTotal() // whole-task total (all agents)
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"workers": stats, "sessions": sessions, "total": tokenTotalDTO(total)})
}
// tokenDailyStats returns global token consumption aggregated by calendar day
@@ -3212,6 +3249,21 @@ func (s *Server) finishTaskChat(taskID string, cancel context.CancelCauseFunc) {
s.chatMu.Unlock()
}
// taskChatStatus reports the authoritative state of the task's main-agent turn.
// Activity timestamps are not a reliable proxy because a tool or LLM call may run
// for minutes without emitting an intermediate frame.
func (s *Server) taskChatStatus(w http.ResponseWriter, r *http.Request) {
t := s.m.ResolveTask(r.PathValue("id"))
if t == nil {
writeErr(w, http.StatusNotFound, "task not found")
return
}
s.chatMu.Lock()
running := s.chatBusy[t.ID]
s.chatMu.Unlock()
writeJSON(w, http.StatusOK, map[string]bool{"running": running})
}
// stopChat aborts the in-flight main-agent turn for a task (manual stop button).
// Mirrors pgStopConversation: cancels only the current turn; planner/workers are
// unaffected and continue running. The user can send a new message immediately.
server/server_mgmt.go
+41 -3
View File
@@ -77,6 +77,44 @@ func canonicalTaskID(raw string) (string, bool) {
return strconv.FormatInt(n, 10), true
}
// beginTaskDelete serializes the delete barrier with pause, resume, admission,
// and FIFO reconciliation. Every lifecycle path rechecks deleting after taking
// concMu, so none can commit a contradictory state once this returns.
func (s *Server) beginTaskDelete(taskID string) bool {
s.concMu.Lock()
defer s.concMu.Unlock()
if s.engine.IsDeleting(taskID) {
return false
}
return s.engine.BeginDelete(taskID)
}
// abortTaskDelete restores the execution barrier from the committed task row,
// not from the in-memory state observed before deletion began. Keeping the task
// paused on an unreadable row is conservative: a later explicit resume can
// safely recover it without allowing work to escape an uncertain delete.
func (s *Server) abortTaskDelete(taskID string) {
s.concMu.Lock()
defer s.concMu.Unlock()
keepPaused := true
if id, err := strconv.ParseInt(taskID, 10, 64); err == nil && s.m != nil && s.m.pg != nil {
if persisted, getErr := s.m.pg.GetTask(id); getErr == nil && persisted != nil {
keepPaused = persisted.Paused || persisted.Queued
} else if task, ok := s.m.Task(taskID); ok {
state := task.lifecycleSnapshot()
keepPaused = state.Paused || state.Queued
if getErr != nil {
log.Printf("[task-delete] task %s 读取持久状态失败，使用内存状态恢复屏障: %v", taskID, getErr)
}
} else if getErr == nil {
// The request targeted a task that does not exist. Do not retain a
// synthetic pause entry after releasing its temporary delete barrier.
keepPaused = false
}
}
s.engine.AbortDelete(taskID, keepPaused)
}
func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
id, ok := canonicalTaskID(r.PathValue("id"))
if !ok {
@@ -88,14 +126,14 @@ func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "invalid JSON: "+err.Error())
return
}
if !s.engine.BeginDelete(id) {
if !s.beginTaskDelete(id) {
writeErr(w, http.StatusConflict, "任务正在删除")
return
}
deleted := false
defer func() {
if !deleted {
s.engine.AbortDelete(id)
s.abortTaskDelete(id)
}
}()
@@ -1754,7 +1792,7 @@ func (s *Server) restoreTasksAfterProfileDelete(pg *db.DB) {
// or clear the explicit chain and restore its Agent/global fallback.
// Resume only intents that were blocked by the exhausted chain; a
// paused task remains paused and terminal tasks remain immutable.
if !isTerminalStatus(task.Status) && s.taskRuntimeAvailable(task, "planner", "worker") {
if !isTerminalStatus(task.lifecycleSnapshot().Status) && s.taskRuntimeAvailable(task, "planner", "worker") {
if _, reopenErr := task.Store.ReopenIntentsByBlockedReason(db.IntentBlockedLLMQuota); reopenErr != nil {
log.Printf("[llm-profile] task %s reopen quota-blocked intents after profile delete: %v", task.ID, reopenErr)
}
server/task_admission_test.go
+790
View File
@@ -0,0 +1,790 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
)
func newAdmissionTestServer(m *Manager, available func(*Task) bool) *Server {
ctx, cancel := context.WithCancel(context.Background())
cancel()
e := NewEngine(m)
planner, worker := new(agent.Planner), new(agent.Worker)
e.SetAuthoritativeAgentResolver(func(task *Task) (*agent.Planner, *agent.Worker) {
if available == nil || available(task) {
return planner, worker
}
return nil, nil
})
return &Server{m: m, engine: e, ctx: ctx, chatBusy: map[string]bool{}, chatCancel: map[string]context.CancelCauseFunc{}}
}
func restoreConcurrencySetting(t *testing.T, m *Manager) func() {
t.Helper()
enabled, limit := m.ConcurrencyLimit()
if limit < 1 {
limit = defaultConcurrencyLimit
}
return func() {
if err := m.SetConcurrency(enabled, limit); err != nil {
t.Errorf("restore concurrency setting: %v", err)
}
}
}
func TestAdmitAlreadyRunningTaskDoesNotQueueAfterLimitDecrease(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
oldEnabled, oldLimit := m.ConcurrencyLimit()
if oldLimit < 1 {
oldLimit = defaultConcurrencyLimit
}
defer m.SetConcurrency(oldEnabled, oldLimit)
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
running, err := m.CreateTask("already admitted", "accept follow-up", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("other admitted task", "occupy lowered limit", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(running.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
// Simulate a task whose long-lived loops were admitted before the limit was
// lowered. The other live task makes the current count exceed the new limit.
s.engine.started.Store(running.ID, true)
queued, err := s.admitTask(running, "resume")
if err != nil {
t.Fatal(err)
}
if queued || running.Queued {
t.Fatalf("already-running task was re-queued: queued=%v task=%+v", queued, running)
}
}
func TestAdmissionDoesNotOverwriteConcurrentTerminalStatus(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(false, defaultConcurrencyLimit); err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("stale admission", "preserve terminal writer", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
// Model an Engine terminal commit after admission captured its in-memory
// snapshot but before the scheduler UPDATE. Direct DB use intentionally leaves
// the live handle stale at its old status for this deterministic interleaving.
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
if err := m.pg.SetStatus(taskID, "done"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
if _, err := s.admitTask(task, "resume"); err == nil {
t.Fatal("stale admission overwrote a concurrent terminal transition")
}
stored, err := m.pg.GetTask(taskID)
if err != nil {
t.Fatal(err)
}
if stored == nil || stored.Status != "done" || stored.Queued || stored.Paused {
t.Fatalf("terminal state changed by stale admission: %+v", stored)
}
}
func TestAdmitKeepsQueuedBootstrapMode(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
oldEnabled, oldLimit := m.ConcurrencyLimit()
if oldLimit < 1 {
oldLimit = defaultConcurrencyLimit
}
defer m.SetConcurrency(oldEnabled, oldLimit)
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("queue holder", "occupy slot", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
queuedTask, err := m.CreateTask("queued bootstrap", "must decompose goals", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(queuedTask.ID, DeleteTaskOptions{})
}()
if err := m.EnqueueTask(queuedTask.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
firstQueuedAt := queuedTask.QueuedAt
s := newAdmissionTestServer(m, nil)
queued, err := s.admitTask(queuedTask, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || queuedTask.QueueMode != "bootstrap" || queuedTask.QueuedAt != firstQueuedAt {
t.Fatalf("re-admission changed bootstrap FIFO state: queued=%v task=%+v", queued, queuedTask)
}
}
func TestTerminalTaskQueuedByAdmissionKeepsExecutionBarrier(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("terminal follow-up", "must wait", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
intentID, err := task.Store.AddIntent(map[string]any{"summary": "follow-up"}, 10, nil, "human")
if err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
queued, err := s.admitTask(task, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || !task.Queued || task.Status != "running" || !s.engine.IsPaused(task.ID) {
t.Fatalf("queued terminal admission lost gate: queued=%v task=%+v enginePaused=%v", queued, task, s.engine.IsPaused(task.ID))
}
if execCtx := s.engine.execContextFor(context.Background(), task.ID); execCtx.Err() == nil {
t.Fatal("queued terminal task received a live execution context")
}
node, err := task.Store.GetNode(intentID)
if err != nil || node == nil || node.State != "open" {
t.Fatalf("queued intent changed before admission: node=%+v err=%v", node, err)
}
}
func TestTimedOutTaskRevivalResetsClockAndSettlingAcrossFIFO(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("timeout revival holder", "occupy slot", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("timed out follow-up", "run with a fresh clock", nil, 120, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
if _, err := m.StampTaskFirstRun(task.ID); err != nil {
t.Fatal(err)
}
oldClock := task.lifecycleSnapshot()
if oldClock.FirstRunAt == 0 || oldClock.DeadlineAt == 0 {
t.Fatalf("test setup did not stamp the original clock: %+v", oldClock)
}
if err := m.SetTaskStatus(task.ID, "timeout"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
s.engine.settling.Store(task.ID, true)
s.engine.deadline.Store(task.ID, oldClock.DeadlineAt)
s.engine.stamped.Store(task.ID, true)
s.engine.coordStarted.Store(task.ID, true)
queued, err := s.admitTask(task, "resume")
if err != nil {
t.Fatal(err)
}
queuedState := task.lifecycleSnapshot()
if !queued || !queuedState.Queued || queuedState.Status != "running" ||
queuedState.FirstRunAt != 0 || queuedState.DeadlineAt != 0 {
t.Fatalf("timed-out task was not reset before FIFO wait: queued=%v state=%+v", queued, queuedState)
}
if s.engine.isSettling(task.ID) {
t.Fatal("timed-out task retained the old settling gate while queued")
}
if _, ok := s.engine.deadline.Load(task.ID); ok {
t.Fatal("timed-out task retained the old in-memory deadline")
}
if _, ok := s.engine.stamped.Load(task.ID); ok {
t.Fatal("timed-out task retained the old first-run stamp")
}
stored, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if stored == nil || stored.TimeoutSeconds != 120 || stored.FirstRunAt != nil || stored.DeadlineAt != nil {
t.Fatalf("persisted timeout revival clock/config mismatch: %+v", stored)
}
// Release the only slot and promote the revived task. The test server uses a
// cancelled root context so no real Agent call races these state assertions.
s.engine.Pause(holder.ID, agent.AbortPausedByUser)
if err := m.ApplyTaskPause(holder.ID); err != nil {
t.Fatal(err)
}
s.reconcileConcurrency()
runningState := task.lifecycleSnapshot()
if runningState.Queued || runningState.Paused || s.engine.IsPaused(task.ID) || s.engine.isSettling(task.ID) {
t.Fatalf("revived task did not leave FIFO ready to execute: state=%+v enginePaused=%v settling=%v",
runningState, s.engine.IsPaused(task.ID), s.engine.isSettling(task.ID))
}
if execCtx := s.engine.execContextFor(context.Background(), task.ID); execCtx.Err() != nil {
t.Fatalf("revived task did not receive a live execution context: %v", context.Cause(execCtx))
}
if _, ok := s.engine.coordStarted.Load(task.ID); !ok {
t.Fatal("revived timeout task did not start a fresh deadline coordinator")
}
}
func TestQueuedTaskCanPauseAndResumeAtFIFOTail(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("pause queued", "resume later", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
if err := m.SetTaskPaused(task.ID, true); err != nil {
t.Fatal(err)
}
s.engine.Pause(task.ID, agent.AbortPausedByUser)
first, err := s.applyTaskControl(task, "resume")
if err != nil || !first.Queued || task.QueuedAt == 0 {
t.Fatalf("first resume: result=%+v task=%+v err=%v", first, task, err)
}
firstQueuedAt := task.QueuedAt
if _, err := s.applyTaskControl(task, "pause"); err != nil {
t.Fatal(err)
}
if !task.Paused || task.Queued || task.QueuedAt != 0 {
t.Fatalf("pause did not leave queue: %+v", task)
}
time.Sleep(2 * time.Millisecond)
second, err := s.applyTaskControl(task, "resume")
if err != nil || !second.Queued {
t.Fatalf("second resume: result=%+v task=%+v err=%v", second, task, err)
}
if task.QueuedAt <= firstQueuedAt {
t.Fatalf("resumed task did not move to FIFO tail: first=%d second=%d", firstQueuedAt, task.QueuedAt)
}
}
func TestReadyFIFOIsAdmittedBeforeNewTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 2); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
older, err := m.CreateTask("older queued", "first", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
incoming, err := m.CreateTask("new admission", "second", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(older.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(incoming.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
if err := m.EnqueueTask(older.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(older.ID, agent.AbortPausedByOrchestrator)
queued, err := s.admitTask(incoming, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || !incoming.Queued || incoming.QueuedAt <= older.QueuedAt {
t.Fatalf("new task bypassed ready FIFO: queued=%v older=%+v incoming=%+v", queued, older, incoming)
}
s.reconcileConcurrency()
if older.Queued || !incoming.Queued {
t.Fatalf("FIFO promotion order wrong: older=%+v incoming=%+v", older, incoming)
}
}
func TestUnavailableTaskReleasesSlotForReadyQueue(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
unavailable, err := m.CreateTask("unavailable llm", "park", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
ready, err := m.CreateTask("ready queued", "run", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(unavailable.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(ready.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, func(task *Task) bool { return task.ID != unavailable.ID })
s.engine.started.Store(unavailable.ID, true)
if err := m.EnqueueTask(ready.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(ready.ID, agent.AbortPausedByOrchestrator)
s.reconcileConcurrency()
if !unavailable.Queued || !s.engine.IsPaused(unavailable.ID) {
t.Fatalf("unavailable task still occupies slot: %+v", unavailable)
}
if ready.Queued || s.engine.IsPaused(ready.ID) {
t.Fatalf("ready task was not promoted: %+v", ready)
}
}
func TestUnavailableTaskWaitsForActiveLLMCallBeforeParking(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
active, err := m.CreateTask("active profile edit", "finish current call", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
waiting, err := m.CreateTask("waiting task", "respect active call", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(active.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(waiting.ID, DeleteTaskOptions{})
}()
available := map[string]bool{active.ID: true, waiting.ID: true}
s := newAdmissionTestServer(m, func(task *Task) bool { return available[task.ID] })
s.engine.started.Store(active.ID, true)
if err := m.EnqueueTask(waiting.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(waiting.ID, agent.AbortPausedByOrchestrator)
// Model a manual chain edit that makes the task's next-call resolution
// unavailable while its old provider is still serving the current call.
s.engine.BeginLLMCall(active.ID)
available[active.ID] = false
s.reconcileConcurrency()
if active.Queued || s.engine.IsPaused(active.ID) {
t.Fatalf("active LLM call was cancelled/parked after profile edit: %+v", active)
}
if !waiting.Queued {
t.Fatalf("waiting task bypassed active provider call: %+v", waiting)
}
s.engine.EndLLMCall(active.ID)
s.reconcileConcurrency()
if !active.Queued || !s.engine.IsPaused(active.ID) {
t.Fatalf("unavailable task was not parked after its call ended: %+v", active)
}
if waiting.Queued || s.engine.IsPaused(waiting.ID) {
t.Fatalf("ready FIFO task was not admitted after active call ended: %+v", waiting)
}
}
func TestRerunAdmissionFailureRestoresIntentState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("failed rerun admission", "restore intent", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer func() {
_, _ = m.pg.Exec(`UPDATE tasks SET deleted_at=NULL WHERE id=$1`, taskID)
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
intentID, err := task.Store.AddIntent(map[string]any{"summary": "retry"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if err := task.Store.SetIntentBlockedReason(intentID, "network_unavailable"); err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
if _, err := m.pg.Exec(`UPDATE tasks SET deleted_at=now() WHERE id=$1`, taskID); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
req := httptest.NewRequest(http.MethodPost, "/api/tasks/1/intents/1/rerun", nil)
req.SetPathValue("id", task.ID)
req.SetPathValue("iid", fmt.Sprintf("%d", intentID))
rec := httptest.NewRecorder()
s.rerunIntent(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
node, err := task.Store.GetNode(intentID)
if err != nil || node == nil || node.State != "blocked" || node.BlockedReason != "network_unavailable" {
t.Fatalf("rerun rollback: node=%+v err=%v", node, err)
}
}
func TestTaskLLMResolutionRejectsInvalidExplicitProfile(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
profileID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("invalid-resolution-%d", time.Now().UnixNano()),
Format: "openai",
Model: "invalid-model",
APIKey: "test-key",
Proxy: "invalid-proxy-without-scheme",
})
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("invalid resolution", "report unavailable", db.TaskCreateOptions{
LLMProfileIDs: []int64{profileID},
})
if err != nil {
_ = m.pg.DeleteProfile(profileID)
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
_ = m.pg.DeleteProfile(profileID)
}()
s := &Server{m: m, provByProfile: map[int64]*provEntry{}}
got, err := s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if got.Available || got.Source != "task_chain" || got.ProfileID == nil || *got.ProfileID != profileID {
t.Fatalf("resolution=%+v, want unavailable explicit task profile", got)
}
}
func TestTaskLLMResolutionReportsDatabaseFailure(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("resolution storage failure", "return server error", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
s := &Server{m: m}
req := httptest.NewRequest(http.MethodGet, "/api/tasks/1/llm/resolution", nil)
req.SetPathValue("id", task.ID)
rec := httptest.NewRecorder()
s.taskLLMResolutionHandler(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestTokenStatsWithoutTaskReturnsStableShape(t *testing.T) {
s := &Server{m: &Manager{tasks: map[string]*Task{}}}
req := httptest.NewRequest(http.MethodGet, "/api/exploration/tokens?task=missing", nil)
rec := httptest.NewRecorder()
s.tokenStats(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
var payload struct {
Workers []db.TokenUsage `json:"workers"`
Sessions []db.SessionTokenUsage `json:"sessions"`
Total TokenTotalDTO `json:"total"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
t.Fatal(err)
}
if payload.Workers == nil || payload.Sessions == nil {
t.Fatalf("empty collections must serialize as arrays: %s", rec.Body.String())
}
}
func TestBatchControlLimitCountsDeduplicatedIDs(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
s := newAdmissionTestServer(m, nil)
request := func(handler http.HandlerFunc, path string, pathValues map[string]string, payload any) *httptest.ResponseRecorder {
body, marshalErr := json.Marshal(payload)
if marshalErr != nil {
t.Fatal(marshalErr)
}
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body))
for key, value := range pathValues {
req.SetPathValue(key, value)
}
rec := httptest.NewRecorder()
handler(rec, req)
return rec
}
rec := request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": []string{" 0009223372036854775807 ", "9223372036854775807", " bad-id ", "bad-id"},
"action": "pause",
})
if rec.Code != http.StatusOK {
t.Fatalf("normalized task ids status=%d body=%s", rec.Code, rec.Body.String())
}
var normalizedPayload struct {
Items []batchControlItem `json:"items"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &normalizedPayload); err != nil {
t.Fatal(err)
}
if len(normalizedPayload.Items) != 2 || normalizedPayload.Items[0].ID != "9223372036854775807" ||
normalizedPayload.Items[1].ID != "bad-id" || normalizedPayload.Items[1].Error != "bad task id" {
t.Fatalf("normalized task response=%+v", normalizedPayload.Items)
}
repeatedTasks := make([]string, maxBatchControlIDs+1)
for i := range repeatedTasks {
repeatedTasks[i] = "missing"
}
rec = request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": repeatedTasks,
"action": "pause",
})
if rec.Code != http.StatusOK {
t.Fatalf("deduplicated task ids status=%d body=%s", rec.Code, rec.Body.String())
}
var taskPayload struct {
Items []batchControlItem `json:"items"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &taskPayload); err != nil {
t.Fatal(err)
}
if len(taskPayload.Items) != 1 || taskPayload.Items[0].ID != "missing" {
t.Fatalf("deduplicated task response=%+v", taskPayload.Items)
}
uniqueTasks := make([]string, maxBatchControlIDs+1)
for i := range uniqueTasks {
uniqueTasks[i] = fmt.Sprintf("missing-%d", i)
}
rec = request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": uniqueTasks,
"action": "pause",
})
if rec.Code != http.StatusBadRequest {
t.Fatalf("unique task ids status=%d body=%s", rec.Code, rec.Body.String())
}
task, err := m.CreateTask("batch intent limit", "deduplicate intent ids", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
repeatedIntents := make([]string, maxBatchControlIDs+1)
for i := range repeatedIntents {
repeatedIntents[i] = "999999"
}
rec = request(s.controlIntentsBatch, "/api/tasks/1/intents/control/batch", map[string]string{"id": task.ID}, map[string]any{
"intent_ids": repeatedIntents,
"action": "pause",
})
if rec.Code != http.StatusOK {
t.Fatalf("deduplicated intent ids status=%d body=%s", rec.Code, rec.Body.String())
}
var intentPayload struct {
Items []batchControlItem `json:"items"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &intentPayload); err != nil {
t.Fatal(err)
}
if len(intentPayload.Items) != 1 || intentPayload.Items[0].ID != "999999" {
t.Fatalf("deduplicated intent response=%+v", intentPayload.Items)
}
uniqueIntents := make([]string, maxBatchControlIDs+1)
for i := range uniqueIntents {
uniqueIntents[i] = strconv.Itoa(i + 1)
}
rec = request(s.controlIntentsBatch, "/api/tasks/1/intents/control/batch", map[string]string{"id": task.ID}, map[string]any{
"intent_ids": uniqueIntents,
"action": "resume",
})
if rec.Code != http.StatusBadRequest {
t.Fatalf("unique intent ids status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestPauseTaskToolPersistsAndDequeuesTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("orchestrated pause", "persist queued pause", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
if err := m.EnqueueTask(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
s.engine.Pause(task.ID, agent.Causef("queued_for_admission", "queued", "test queue barrier"))
input, err := json.Marshal(map[string]string{"task_id": task.ID})
if err != nil {
t.Fatal(err)
}
if _, err := s.toolPauseTask().Call(context.Background(), input, nil); err != nil {
t.Fatal(err)
}
if !task.Paused || task.Queued || task.QueueMode != "bootstrap" || !s.engine.IsPaused(task.ID) {
t.Fatalf("tool pause did not persist/dequeue task: %+v enginePaused=%v", task, s.engine.IsPaused(task.ID))
}
stored, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if stored == nil || !stored.Paused || stored.Queued {
t.Fatalf("stored task=%+v, want paused and dequeued", stored)
}
}
func TestPauseTaskToolUsesOrchestratorCancellationCause(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("orchestrated pause cause", "keep audit reason", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := newAdmissionTestServer(m, nil)
execCtx := s.engine.execContextFor(context.Background(), task.ID)
input, err := json.Marshal(map[string]string{"task_id": task.ID})
if err != nil {
t.Fatal(err)
}
if _, err := s.toolPauseTask().Call(context.Background(), input, nil); err != nil {
t.Fatal(err)
}
if code, _, _, ok := agent.AbortReason(execCtx); !ok || code != "paused_by_orchestrator" {
t.Fatalf("pause_task cancellation code=%q ok=%v, want paused_by_orchestrator", code, ok)
}
}
server/task_control.go
+314
View File
@@ -0,0 +1,314 @@
package server
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
)
const maxBatchControlIDs = 100
type taskControlResult struct {
ID string `json:"id"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
Status string `json:"status"`
}
type intentControlResult struct {
ID int64 `json:"id"`
State string `json:"state"`
Deleted *db.IntentCleanup `json:"deleted,omitempty"`
}
type batchControlItem struct {
ID string `json:"id"`
OK bool `json:"ok"`
Status string `json:"status,omitempty"`
State string `json:"state,omitempty"`
Queued bool `json:"queued,omitempty"`
Error string `json:"error,omitempty"`
}
func (s *Server) resumeAdmissionMode(t *Task) string {
if t == nil {
return "resume"
}
lifecycle := t.lifecycleSnapshot()
if lifecycle.QueueMode == "bootstrap" {
return "bootstrap"
}
if lifecycle.FirstRunAt == 0 {
goals, err := t.Store.ListByKind(db.KindGoal, 1)
if err == nil && len(goals) == 0 {
return "bootstrap"
}
}
return "resume"
}
// applyTaskControl is shared by the single and batch endpoints. Task resume is
// deliberately limited to paused tasks; reruns and finding follow-ups use
// admitTask directly when they need to revive a terminal task.
func (s *Server) applyTaskControl(t *Task, action string) (taskControlResult, error) {
return s.applyTaskControlWithCause(t, action, agent.AbortPausedByUser)
}
func (s *Server) applyTaskControlWithCause(t *Task, action string, pauseCause error) (taskControlResult, error) {
if t == nil {
return taskControlResult{}, fmt.Errorf("task not found")
}
out := taskControlResult{ID: t.ID}
switch action {
case "pause":
s.concMu.Lock()
defer s.concMu.Unlock()
current, exists := s.m.Task(t.ID)
if !exists || current != t || s.engine.IsDeleting(t.ID) {
return out, fmt.Errorf("任务正在删除，无法控制")
}
if !s.engine.beginTaskOperation(t.ID) {
return out, fmt.Errorf("任务正在删除，无法控制")
}
defer s.engine.decInflight(t.ID)
lifecycle := t.lifecycleSnapshot()
if isTerminalStatus(lifecycle.Status) {
return out, fmt.Errorf("终态任务不能执行暂停")
}
if lifecycle.Paused {
return out, fmt.Errorf("任务已经暂停")
}
wasQueued := lifecycle.Queued
wasEnginePaused := s.engine.IsPaused(t.ID)
if pauseCause == nil {
pauseCause = agent.AbortPausedByUser
}
s.engine.Pause(t.ID, pauseCause)
if err := s.m.ApplyTaskPause(t.ID); err != nil {
if !wasEnginePaused && !wasQueued {
s.engine.Resume(t)
}
return out, err
}
// Main Agent is independently cancellable. Only cancel its current turn
// after the persistent pause commits, so a failed control request is fully
// compensated and does not lose an otherwise valid conversation turn.
s.cancelTaskChat(t.ID, agent.AbortChatPausedWithTask)
out.Paused, out.Status = true, "paused"
go s.reconcileConcurrency()
case "resume":
queued, err := s.admitPausedTask(t)
if err != nil {
return out, err
}
out.Queued = queued
out.Status = map[bool]string{true: "queued", false: "running"}[queued]
default:
return out, fmt.Errorf("action must be pause|resume")
}
log.Printf("[task] #%s %s", t.ID, map[string]string{"pause": "已暂停", "resume": "已继续"}[action])
return out, nil
}
func (s *Server) applyIntentControl(ctx context.Context, t *Task, iid int64, action string) (intentControlResult, error) {
out := intentControlResult{ID: iid}
node, err := t.Store.GetNode(iid)
if err != nil {
return out, err
}
if node == nil {
if inherited, sourceErr := t.Store.GetNodeWithSources(iid); sourceErr == nil && inherited != nil && inherited.Inherited {
return out, fmt.Errorf("继承意图为只读，不能控制")
}
return out, fmt.Errorf("intent not found")
}
if node.Kind != db.KindIntent {
return out, fmt.Errorf("node is not an intent")
}
switch action {
case "pause":
if node.State != "running" {
return out, fmt.Errorf("仅运行中的意图可以暂停")
}
if err := s.engine.ControlWork(ctx, iid, "pause"); err != nil {
return out, err
}
out.State = "paused"
case "resume":
if node.State != "paused" {
return out, fmt.Errorf("仅已暂停的意图可以恢复")
}
changed, err := t.Store.CompareAndSetIntentState(iid, "paused", "open")
if err != nil {
return out, err
}
if !changed {
return out, fmt.Errorf("%w: 意图不再是 paused 状态", db.ErrIntentStateConflict)
}
t.Notify()
out.State = "open"
case "cancel":
if node.State != "running" && node.State != "paused" {
return out, fmt.Errorf("仅运行中或已暂停的意图可以取消")
}
if node.State == "running" {
if err := s.engine.ControlWork(ctx, iid, "cancel"); err != nil {
return out, err
}
}
deleted, err := t.Store.CancelIntent(iid)
if err != nil {
return out, err
}
t.Notify()
out.State, out.Deleted = "cancelled", &deleted
default:
return out, fmt.Errorf("action must be pause|resume|cancel")
}
return out, nil
}
func (s *Server) controlTasksBatch(w http.ResponseWriter, r *http.Request) {
r.Body = http.MaxBytesReader(w, r.Body, 256<<10)
var req struct {
TaskIDs []string `json:"task_ids"`
Action string `json:"action"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "bad json: "+err.Error())
return
}
if req.Action != "pause" && req.Action != "resume" {
writeErr(w, 400, "action must be pause|resume")
return
}
type parsedTaskID struct {
id string
valid bool
}
seen := map[string]bool{}
seenInvalid := map[string]bool{}
taskIDs := make([]parsedTaskID, 0, len(req.TaskIDs))
for _, raw := range req.TaskIDs {
trimmed := strings.TrimSpace(raw)
id, valid := canonicalTaskID(trimmed)
if !valid {
if seenInvalid[trimmed] {
continue
}
seenInvalid[trimmed] = true
taskIDs = append(taskIDs, parsedTaskID{id: trimmed})
continue
}
if seen[id] {
continue
}
seen[id] = true
taskIDs = append(taskIDs, parsedTaskID{id: id, valid: true})
}
if len(taskIDs) == 0 || len(taskIDs) > maxBatchControlIDs {
writeErr(w, 400, fmt.Sprintf("task_ids 数量必须为 1-%d", maxBatchControlIDs))
return
}
items := make([]batchControlItem, 0, len(taskIDs))
for _, parsed := range taskIDs {
item := batchControlItem{ID: parsed.id}
if !parsed.valid {
item.Error = "bad task id"
items = append(items, item)
continue
}
t, ok := s.m.Task(parsed.id)
if !ok {
item.Error = "task not found"
items = append(items, item)
continue
}
result, err := s.applyTaskControl(t, req.Action)
if err != nil {
item.Error = err.Error()
} else {
item.OK, item.Status, item.Queued = true, result.Status, result.Queued
}
items = append(items, item)
}
writeJSON(w, 200, map[string]any{"items": items})
}
func (s *Server) controlIntentsBatch(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
r.Body = http.MaxBytesReader(w, r.Body, 256<<10)
var req struct {
IntentIDs []string `json:"intent_ids"`
Action string `json:"action"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "bad json: "+err.Error())
return
}
if req.Action != "pause" && req.Action != "resume" {
writeErr(w, 400, "action must be pause|resume")
return
}
type parsedIntentID struct {
raw string
id int64
valid bool
}
seen := map[int64]bool{}
seenInvalid := map[string]bool{}
intentIDs := make([]parsedIntentID, 0, len(req.IntentIDs))
for _, raw := range req.IntentIDs {
iid, err := strconv.ParseInt(raw, 10, 64)
if err != nil || iid <= 0 {
if seenInvalid[raw] {
continue
}
seenInvalid[raw] = true
intentIDs = append(intentIDs, parsedIntentID{raw: raw})
continue
}
if seen[iid] {
continue
}
seen[iid] = true
intentIDs = append(intentIDs, parsedIntentID{raw: raw, id: iid, valid: true})
}
if len(intentIDs) == 0 || len(intentIDs) > maxBatchControlIDs {
writeErr(w, 400, fmt.Sprintf("intent_ids 数量必须为 1-%d", maxBatchControlIDs))
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除，无法控制意图")
return
}
defer s.engine.decInflight(t.ID)
items := make([]batchControlItem, 0, len(intentIDs))
for _, parsed := range intentIDs {
item := batchControlItem{ID: parsed.raw}
if !parsed.valid {
item.Error = "bad intent id"
items = append(items, item)
continue
}
result, err := s.applyIntentControl(r.Context(), t, parsed.id, req.Action)
if err != nil {
item.Error = err.Error()
} else {
item.OK, item.State = true, result.State
}
items = append(items, item)
}
writeJSON(w, 200, map[string]any{"items": items})
}
server/task_delete_barrier_test.go
+146 -2
View File
@@ -42,7 +42,7 @@ func TestEngineDeleteBarrierRejectsNewTaskOperations(t *testing.T) {
if e.BeginDelete(taskID) {
t.Fatal("duplicate delete must not acquire the barrier")
}
e.AbortDelete(taskID)
e.AbortDelete(taskID, false)
if e.IsPaused(taskID) {
t.Fatal("aborted delete should restore a previously running task")
}
@@ -58,12 +58,130 @@ func TestAbortDeletePreservesExistingPause(t *testing.T) {
if !e.BeginDelete("42") {
t.Fatal("delete barrier should be installed")
}
e.AbortDelete("42")
e.AbortDelete("42", true)
if !e.IsPaused("42") {
t.Fatal("aborted delete must preserve a pre-existing user pause")
}
}
func TestBeginTaskDeleteSerializesWithLifecycleLock(t *testing.T) {
s := &Server{engine: NewEngine(nil)}
s.concMu.Lock()
done := make(chan bool, 1)
go func() { done <- s.beginTaskDelete("42") }()
select {
case <-done:
t.Fatal("delete barrier bypassed the lifecycle lock")
case <-time.After(20 * time.Millisecond):
}
if s.engine.IsDeleting("42") {
t.Fatal("delete barrier was installed before concMu was released")
}
s.concMu.Unlock()
select {
case ok := <-done:
if !ok {
t.Fatal("delete barrier was not installed after lifecycle lock release")
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for delete barrier")
}
s.engine.AbortDelete("42", false)
}
func TestDeletingRunningTaskKeepsConcurrencySlotUntilCommit(t *testing.T) {
task := &Task{ID: "42", Status: "running"}
s := &Server{engine: NewEngine(nil)}
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install delete barrier")
}
if !s.occupiesConcurrencySlot(task) {
t.Fatal("temporary delete pause released the running task's slot before commit")
}
s.engine.AbortDelete(task.ID, false)
}
func TestTaskLifecycleRechecksDeleteBarrierAfterConcLock(t *testing.T) {
task := &Task{ID: "42", notify: make(chan struct{}, 1)}
m := &Manager{tasks: map[string]*Task{task.ID: task}}
s := &Server{m: m, engine: NewEngine(m)}
s.concMu.Lock()
started := make(chan struct{})
pauseDone := make(chan error, 1)
go func() {
close(started)
_, err := s.applyTaskControl(task, "pause")
pauseDone <- err
}()
<-started
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install delete barrier")
}
s.concMu.Unlock()
if err := <-pauseDone; err == nil || !strings.Contains(err.Error(), "正在删除") {
t.Fatalf("pause passed a delete barrier installed while waiting for concMu: %v", err)
}
s.engine.AbortDelete(task.ID, false)
s.concMu.Lock()
started = make(chan struct{})
admitDone := make(chan error, 1)
go func() {
close(started)
_, err := s.admitTask(task, "resume")
admitDone <- err
}()
<-started
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install second delete barrier")
}
s.concMu.Unlock()
if err := <-admitDone; err == nil || !strings.Contains(err.Error(), "being deleted") {
t.Fatalf("admission passed a delete barrier installed while waiting for concMu: %v", err)
}
s.engine.AbortDelete(task.ID, false)
}
func TestAbortTaskDeleteUsesPersistedPauseAndQueueState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("delete abort state", "restore persisted barrier", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := &Server{m: m, engine: NewEngine(m)}
// An engine-only pause is not authoritative. The committed task is running,
// so aborting deletion must clear this stale in-memory barrier.
s.engine.Pause(task.ID, agent.AbortPausedByUser)
if !s.beginTaskDelete(task.ID) {
t.Fatal("install first delete barrier")
}
s.abortTaskDelete(task.ID)
if s.engine.IsPaused(task.ID) {
t.Fatal("stale engine pause survived although the persisted task is running")
}
// A queued row must retain its execution barrier even when BeginDelete saw the
// engine as running. This specifically guards against restoring a stale snapshot.
if err := m.EnqueueTask(task.ID, "resume"); err != nil {
t.Fatal(err)
}
if !s.beginTaskDelete(task.ID) {
t.Fatal("install second delete barrier")
}
s.abortTaskDelete(task.ID)
if !s.engine.IsPaused(task.ID) {
t.Fatal("persisted queued task lost its execution barrier after abort")
}
}
func TestStopTaskCancelsRuntimeAndClearsLifecycleState(t *testing.T) {
e := NewEngine(nil)
const taskID = "42"
@@ -182,6 +300,32 @@ func TestWaitTaskQuiescentHonorsContext(t *testing.T) {
}
}
func TestTaskChatStatusUsesAuthoritativeBusyState(t *testing.T) {
m := &Manager{tasks: map[string]*Task{"7": {ID: "7"}}}
s := &Server{
m: m,
chatBusy: map[string]bool{"7": true},
chatCancel: map[string]context.CancelCauseFunc{},
}
req := httptest.NewRequest(http.MethodGet, "/api/tasks/7/chat/status", nil)
req.SetPathValue("id", "7")
rec := httptest.NewRecorder()
s.taskChatStatus(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String())
}
var got struct {
Running bool `json:"running"`
}
if err := json.NewDecoder(rec.Body).Decode(&got); err != nil {
t.Fatal(err)
}
if !got.Running {
t.Fatal("busy main-agent turn was reported idle")
}
}
func TestTaskChatUploadRejectsDeleteBarrierBeforeCreatingDirectory(t *testing.T) {
dataDir := t.TempDir()
m := &Manager{dir: dataDir, tasks: map[string]*Task{"7": {ID: "7"}}}
server/task_llm.go
+1 -1
View File
@@ -375,7 +375,7 @@ func (s *Server) syncTaskLLMState(pt *db.Task) {
id := fmt.Sprintf("%d", pt.ID)
s.m.mu.Lock()
if task := s.m.tasks[id]; task != nil {
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason)
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMChainRevision, pt.LLMFailoverState, pt.LLMFailoverReason)
}
s.m.mu.Unlock()
}
server/task_llm_test.go
+20 -2
View File
@@ -122,7 +122,7 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
defer wg.Done()
for i := 0; i < 500; i++ {
id := int64(offset*1000 + i + 1)
task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "")
task.setLLMState(&id, &id, []int64{id, id + 1}, int64(i), "ready", "")
}
}(writer)
}
@@ -146,6 +146,24 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
}
}
func TestTaskLLMStateRejectsOlderRevision(t *testing.T) {
t.Parallel()
task := &Task{ID: "7"}
current := int64(22)
stale := int64(11)
if !task.setLLMState(&current, &current, []int64{current}, 2, "ready", "") {
t.Fatal("initial LLM state was rejected")
}
if task.setLLMState(&stale, &stale, []int64{stale}, 1, "ready", "stale") {
t.Fatal("older LLM snapshot was accepted")
}
state := task.llmStateSnapshot()
if state.ChainRevision != 2 || state.ActiveID == nil || *state.ActiveID != current ||
len(state.ProfileIDs) != 1 || state.ProfileIDs[0] != current || state.FailoverReason != "" {
t.Fatalf("older LLM snapshot overwrote current state: %+v", state)
}
}
func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) {
t.Parallel()
quota := errors.New("openai: status 402: insufficient_quota")
@@ -381,7 +399,7 @@ func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) {
engine.started.Store("7", true)
task := &Task{ID: "7", Status: "running"}
profileID := int64(11)
task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted")
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
s := &Server{engine: engine}
if got := s.resolvedTaskStatus(task); got != "running" {
t.Fatalf("exhausted started status=%q, want running", got)
server/task_resolution.go
+123
View File
@@ -0,0 +1,123 @@
package server
import (
"fmt"
"net/http"
"github.com/Autumn-27/artex/db"
)
type taskLLMResolution struct {
ProfileID *int64 `json:"profile_id,omitempty"`
Name string `json:"name"`
Format string `json:"format"`
Model string `json:"model"`
Source string `json:"source"`
Available bool `json:"available"`
Reason string `json:"reason,omitempty"`
}
func (s *Server) resolutionFromProfile(p *db.LLMProfile, source string) taskLLMResolution {
if p == nil {
return taskLLMResolution{Source: source, Reason: "LLM 配置不存在"}
}
id := p.ID
result := taskLLMResolution{
ProfileID: &id,
Name: p.Name,
Format: p.Format,
Model: p.Model,
Source: source,
}
if p.APIKey == "" {
result.Reason = "LLM 配置未设置 API Key"
return result
}
if _, _, ok := s.providerForProfile(p.ID); !ok {
result.Reason = "LLM 配置格式或参数无效"
return result
}
result.Available = true
return result
}
// resolveTaskRoleLLM mirrors taskLLMRuntime.current without creating a provider:
// explicit task chain -> role Agent binding -> active/global environment config.
// Database failures are returned rather than represented as an unavailable model.
func (s *Server) resolveTaskRoleLLM(t *Task, agentKey string) (taskLLMResolution, error) {
if s == nil || s.m == nil || s.m.pg == nil {
return taskLLMResolution{}, fmt.Errorf("database unavailable")
}
state := t.llmStateSnapshot()
if len(state.ProfileIDs) > 0 {
if state.ActiveID == nil {
return taskLLMResolution{Source: "task_chain", Reason: "任务 LLM 配置链额度已耗尽"}, nil
}
p, err := s.m.pg.ProfileByID(*state.ActiveID)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load task LLM profile: %w", err)
}
return s.resolutionFromProfile(p, "task_chain"), nil
}
a, err := s.m.pg.GetAgentByKey(agentKey)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load agent %q: %w", agentKey, err)
}
if a != nil && a.LLMProfileID != nil {
p, err := s.m.pg.ProfileByID(*a.LLMProfileID)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load agent LLM profile: %w", err)
}
if p != nil {
resolved := s.resolutionFromProfile(p, "agent_binding")
if resolved.Available {
return resolved, nil
}
}
}
p, err := s.m.pg.ActiveProfile()
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load global LLM profile: %w", err)
}
if p != nil {
resolved := s.resolutionFromProfile(p, "global_profile")
if resolved.Available {
return resolved, nil
}
}
s.cfgMu.Lock()
on, name, cfg := s.llmOn, s.llmProf, s.llmCfg
s.cfgMu.Unlock()
if on {
if name == "" {
name = "全局配置"
}
return taskLLMResolution{
Name: name, Format: cfg.Provider(), Model: cfg.Model,
Source: "environment", Available: true,
}, nil
}
return taskLLMResolution{Source: "global", Reason: "没有可用的 LLM 配置"}, nil
}
func (s *Server) taskLLMResolutionHandler(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
roles := map[string]taskLLMResolution{}
for responseKey, agentKey := range map[string]string{
"mainagent": "mainagent",
"planner": "planner",
"worker": "worker",
} {
resolved, err := s.resolveTaskRoleLLM(t, agentKey)
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
roles[responseKey] = resolved
}
writeJSON(w, http.StatusOK, roles)
}
server/task_templates.go
+172
View File
@@ -0,0 +1,172 @@
package server
import (
"errors"
"fmt"
"net/http"
"strings"
"unicode/utf8"
"github.com/Autumn-27/artex/db"
)
const maxTaskTemplateRequestBytes = 512 << 10
// taskTemplateRequest uses pointers so PATCH can distinguish omitted fields from
// explicit empty values. Empty values are still rejected by the DB validator.
type taskTemplateRequest struct {
Name *string `json:"name"`
Description *string `json:"description"`
Goal *string `json:"goal"`
}
func decodeTaskTemplateRequest(w http.ResponseWriter, r *http.Request, req *taskTemplateRequest) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxTaskTemplateRequestBytes)
if err := decode(r, req); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, err.Error())
}
return false
}
return true
}
func validateTaskTemplateRequest(req taskTemplateRequest) error {
checks := []struct {
name string
value *string
limit int
}{
{name: "name", value: req.Name, limit: db.MaxTaskTemplateNameRunes},
{name: "description", value: req.Description, limit: db.MaxTaskTemplateTextRunes},
{name: "goal", value: req.Goal, limit: db.MaxTaskTemplateTextRunes},
}
for _, check := range checks {
if check.value == nil {
continue
}
value := strings.TrimSpace(*check.value)
if check.name == "name" {
value = strings.Join(strings.Fields(value), " ")
}
if utf8.RuneCountInString(value) > check.limit {
return fmt.Errorf("%s 最多 %d 个字符", check.name, check.limit)
}
}
return nil
}
func writeTaskTemplateErr(w http.ResponseWriter, err error) {
switch {
case errors.Is(err, db.ErrTaskTemplateInvalid):
writeErr(w, http.StatusBadRequest, err.Error())
case errors.Is(err, db.ErrTaskTemplateNameConflict):
writeErr(w, http.StatusConflict, "模板名称已存在")
case errors.Is(err, db.ErrTaskTemplateNotFound):
writeErr(w, http.StatusNotFound, "task template not found")
default:
writeErr(w, http.StatusInternalServerError, err.Error())
}
}
func (s *Server) pgListTaskTemplates(w http.ResponseWriter, _ *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
templates, err := pg.ListTaskTemplates()
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]any{"templates": templates})
}
func (s *Server) pgCreateTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
var req taskTemplateRequest
if !decodeTaskTemplateRequest(w, r, &req) {
return
}
if err := validateTaskTemplateRequest(req); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
template, err := pg.CreateTaskTemplate(db.TaskTemplateInput{
Name: stringValue(req.Name),
Description: stringValue(req.Description),
Goal: stringValue(req.Goal),
})
if err != nil {
writeTaskTemplateErr(w, err)
return
}
writeJSON(w, http.StatusCreated, template)
}
func (s *Server) pgUpdateTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, http.StatusBadRequest, "bad task template id")
return
}
var req taskTemplateRequest
if !decodeTaskTemplateRequest(w, r, &req) {
return
}
if err := validateTaskTemplateRequest(req); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
if req.Name == nil && req.Description == nil && req.Goal == nil {
writeErr(w, http.StatusBadRequest, "至少需要提供 name、description 或 goal")
return
}
template, err := pg.PatchTaskTemplate(id, db.TaskTemplatePatch{
Name: req.Name, Description: req.Description, Goal: req.Goal,
})
if err != nil {
writeTaskTemplateErr(w, err)
return
}
writeJSON(w, http.StatusOK, template)
}
func (s *Server) pgDeleteTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, http.StatusBadRequest, "bad task template id")
return
}
deleted, err := pg.DeleteTaskTemplate(id)
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
if !deleted {
writeErr(w, http.StatusNotFound, "task template not found")
return
}
writeJSON(w, http.StatusOK, map[string]any{"deleted": id})
}
func stringValue(value *string) string {
if value == nil {
return ""
}
return *value
}
server/task_templates_test.go
+163
View File
@@ -0,0 +1,163 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestTaskTemplateHTTPCRUD(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
h := s.Handler()
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
do := func(method, path string, body any) (int, map[string]any) {
var data []byte
if body != nil {
data, _ = json.Marshal(body)
}
req := httptest.NewRequest(method, path, bytes.NewReader(data))
req.Header.Set("Authorization", "Bearer "+token)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
var out map[string]any
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out
}
name := fmt.Sprintf("http-template-%d", time.Now().UnixNano())
code, created := do(http.MethodPost, "/api/task-templates", map[string]any{
"name": name, "description": "description", "goal": "goal",
})
if code != http.StatusCreated {
t.Fatalf("create status=%d body=%v", code, created)
}
id := int64(created["id"].(float64))
t.Cleanup(func() { _, _ = m.pg.DeleteTaskTemplate(id) })
code, duplicate := do(http.MethodPost, "/api/task-templates", map[string]any{
"name": " " + name + " ", "description": "other", "goal": "other",
})
if code != http.StatusConflict {
t.Fatalf("duplicate status=%d body=%v", code, duplicate)
}
code, updated := do(http.MethodPatch, fmt.Sprintf("/api/task-templates/%d", id), map[string]any{"goal": "patched goal"})
if code != http.StatusOK || updated["goal"] != "patched goal" || updated["description"] != "description" {
t.Fatalf("patch status=%d body=%v", code, updated)
}
code, _ = do(http.MethodPatch, fmt.Sprintf("/api/task-templates/%d", id), map[string]any{
"goal": strings.Repeat("界", db.MaxTaskTemplateTextRunes+1),
})
if code != http.StatusBadRequest {
t.Fatalf("overlong patch status=%d, want %d", code, http.StatusBadRequest)
}
code, _ = do(http.MethodPost, "/api/task-templates", map[string]any{
"name": "oversized", "description": strings.Repeat("x", maxTaskTemplateRequestBytes+1), "goal": "goal",
})
if code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized create status=%d, want %d", code, http.StatusRequestEntityTooLarge)
}
code, listed := do(http.MethodGet, "/api/task-templates", nil)
if code != http.StatusOK || listed["templates"] == nil {
t.Fatalf("list status=%d body=%v", code, listed)
}
code, deleted := do(http.MethodDelete, fmt.Sprintf("/api/task-templates/%d", id), nil)
if code != http.StatusOK || int64(deleted["deleted"].(float64)) != id {
t.Fatalf("delete status=%d body=%v", code, deleted)
}
}
func TestConversationPatchReturnsPinState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
conversation, err := m.pg.CreateConversation("mainagent", "pin through http", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = m.pg.DeleteConversation(conversation.ID) })
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
body, _ := json.Marshal(map[string]any{"title": " pinned title ", "pinned": true})
req := httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec := httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("patch status=%d body=%s", rec.Code, rec.Body.String())
}
var updated struct {
Title string `json:"title"`
Pinned bool `json:"pinned"`
PinnedAt *string `json:"pinned_at"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &updated); err != nil {
t.Fatal(err)
}
if updated.Title != "pinned title" || !updated.Pinned || updated.PinnedAt == nil {
t.Fatalf("unexpected conversation patch response: %+v", updated)
}
body, _ = json.Marshal(map[string]any{"title": strings.Repeat("会", maxConversationTitleRunes+1)})
req = httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("overlong conversation title status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{"title": strings.Repeat("x", maxConversationRequestBytes+1)})
req = httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized conversation patch status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{
"agent_key": "mainagent", "title": strings.Repeat("会", maxConversationTitleRunes+1),
})
req = httptest.NewRequest(http.MethodPost, "/api/conversations", bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("overlong conversation create status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{
"agent_key": "mainagent", "title": strings.Repeat("x", maxConversationRequestBytes+1),
})
req = httptest.NewRequest(http.MethodPost, "/api/conversations", bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized conversation create status=%d body=%s", rec.Code, rec.Body.String())
}
}
web/src/app/(main)/chat/page.tsx
+187 -78
View File
@@ -1,20 +1,27 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { Bot, ChevronDownIcon, Loader2Icon, PaperclipIcon, PlusIcon, SendIcon, Square, Trash2Icon, XIcon, ZapIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Textarea } from "@/components/ui/textarea";
import { Badge } from "@/components/ui/badge";
import { ScrollArea } from "@/components/ui/scroll-area";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
Bot,
ChevronDownIcon,
Loader2Icon,
MoreHorizontalIcon,
PaperclipIcon,
PencilIcon,
PinIcon,
PinOffIcon,
PlusIcon,
SendIcon,
Square,
Trash2Icon,
XIcon,
ZapIcon,
} from "lucide-react";
import { toast } from "sonner";
import { TodoPopover } from "@/components/todo-popover";
import { Transcript } from "@/components/transcript";
import {
AlertDialog,
AlertDialogAction,
@@ -24,15 +31,21 @@ import {
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
AlertDialogTrigger,
} from "@/components/ui/alert-dialog";
import { Transcript } from "@/components/transcript";
import { TodoPopover } from "@/components/todo-popover";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
DropdownMenu,
DropdownMenuContent,
DropdownMenuGroup,
DropdownMenuItem,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import { api } from "@/lib/api";
import type { Activity, Agent, ChatAttachment, Conversation, LLMProfile } from "@/lib/types";
import { cn } from "@/lib/utils";
@@ -46,8 +59,8 @@ function fmtBytes(n: number): string {
// fmtTokens renders a compact token count (1234 → 1.2k, 2_000_000 → 2M).
function fmtTokens(n: number): string {
if (n >= 1_000_000) return (n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1) + "M";
if (n >= 1000) return (n / 1000).toFixed(n >= 10000 ? 0 : 1) + "k";
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1)}M`;
if (n >= 1000) return `${(n / 1000).toFixed(n >= 10000 ? 0 : 1)}k`;
return String(n);
}
@@ -66,6 +79,10 @@ function fmtDuration(ms: number): string {
// thread stays snappy (only ~a page of rows is in the DOM until you scroll up).
const HISTORY_PAGE = 200;
function conversationIsPinned(conversation: Conversation): boolean {
return conversation.pinned ?? Boolean(conversation.pinned_at);
}
// LiveBadge is the small pulsing "实时" chip reused from the task's main-agent
// console — shown while a turn is streaming.
function LiveBadge() {
@@ -113,7 +130,7 @@ function Composer({
const atts = attachments ?? [];
function onKeyDown(e: React.KeyboardEvent<HTMLTextAreaElement>) {
// Enter 发送；Shift+Enter 换行。isComposing 时是中文等输入法在选字,不拦截。
if (e.key !== "Enter" || e.shiftKey || e.nativeEvent.isComposing) return;
if (e.key !== "Enter" || e.shiftKey || e.nativeEvent.isComposing || e.nativeEvent.keyCode === 229) return;
e.preventDefault();
onSend();
}
@@ -144,8 +161,8 @@ function Composer({
))}
</div>
)}
<div className="flex items-end gap-2">
{leftSlot}
<div className="flex flex-wrap items-end gap-2">
{leftSlot ? <div className="w-full sm:w-auto">{leftSlot}</div> : null}
{onPickFiles && (
<>
<input
@@ -170,7 +187,7 @@ function Composer({
</>
)}
<Textarea
className="max-h-40 min-h-10 flex-1 resize-none"
className="max-h-40 min-h-10 min-w-0 flex-1 resize-none"
rows={1}
placeholder={placeholder}
value={value}
@@ -234,7 +251,10 @@ function LLMProfileRow({
{/* default option */}
<button
type="button"
onClick={() => { onChange(null); setOpen(false); }}
onClick={() => {
onChange(null);
setOpen(false);
}}
className={cn(
"flex w-full flex-col rounded px-2 py-1.5 text-left hover:bg-accent",
selected == null && "bg-accent",
@@ -242,21 +262,28 @@ function LLMProfileRow({
>
<span className="text-sm">默认{activeDefault ? `（${activeDefault.name}）` : ""}</span>
{activeDefault && (
<span className="text-muted-foreground text-[11px]">{activeDefault.format} · {activeDefault.model}</span>
<span className="text-muted-foreground text-[11px]">
{activeDefault.format} · {activeDefault.model}
</span>
)}
</button>
{profiles.map((p) => (
<button
key={p.id}
type="button"
onClick={() => { onChange(Number(p.id)); setOpen(false); }}
onClick={() => {
onChange(Number(p.id));
setOpen(false);
}}
className={cn(
"flex w-full flex-col rounded px-2 py-1.5 text-left hover:bg-accent",
selected === Number(p.id) && "bg-accent",
)}
>
<span className="text-sm">{p.name}</span>
<span className="text-muted-foreground text-[11px]">{p.format} · {p.model}</span>
<span className="text-muted-foreground text-[11px]">
{p.format} · {p.model}
</span>
</button>
))}
</PopoverContent>
@@ -326,7 +353,7 @@ function DraftChat({
const agentPicker = (
<Select value={agentKey} onValueChange={setAgentKey}>
<SelectTrigger size="sm" className="w-40 shrink-0">
<SelectTrigger size="sm" className="w-full sm:w-40">
<SelectValue placeholder="选择 Agent…" />
</SelectTrigger>
<SelectContent>
@@ -336,7 +363,9 @@ function DraftChat({
<Bot className="size-3.5" />
{a.name}
{!a.builtin && (
<Badge variant="outline" className="px-1 py-0 text-[9px]">自定义</Badge>
<Badge variant="outline" className="px-1 py-0 text-[9px]">
自定义
</Badge>
)}
</span>
</SelectItem>
@@ -353,9 +382,7 @@ function DraftChat({
<Bot className="text-primary size-6" />
</div>
<div className="text-sm font-medium">开始和「{agent?.name ?? "Agent"}」对话</div>
{agent?.description && (
<p className="text-muted-foreground max-w-md text-xs">{agent.description}</p>
)}
{agent?.description && <p className="text-muted-foreground max-w-md text-xs">{agent.description}</p>}
</div>
<Composer
@@ -424,10 +451,7 @@ function ChatView({
}
// conversation-scoped detail fetcher for the reused Transcript renderer.
const fetchDetail = React.useCallback(
(seq: number) => api.conversationMsgDetail(conv.id, seq),
[conv.id],
);
const fetchDetail = React.useCallback((seq: number) => api.conversationMsgDetail(conv.id, seq), [conv.id]);
// seq of the most-recent TodoWrite tool call (for the Todo popover); null if none.
const latestTodoSeq = React.useMemo(() => {
@@ -462,7 +486,9 @@ function ChatView({
setHasMore(r.hasMore);
setRunning(r.running);
})
.catch(() => {});
.catch(() => {
/* The empty state remains usable when history loading fails. */
});
return () => {
live = false;
};
@@ -497,9 +523,7 @@ function ChatView({
const contentRef = React.useRef<HTMLDivElement | null>(null);
const atBottomRef = React.useRef(true);
const viewport = React.useCallback(
() =>
(contentRef.current?.closest('[data-slot="scroll-area-viewport"]') as HTMLElement | null) ??
null,
() => (contentRef.current?.closest('[data-slot="scroll-area-viewport"]') as HTMLElement | null) ?? null,
[],
);
// Scroll-up loads one older page and prepends it, preserving the visual position
@@ -543,6 +567,7 @@ function ChatView({
return () => vp.removeEventListener("scroll", onScroll);
}, [viewport, loadEarlier]);
// open/switch a conversation → jump to the latest (bottom)
// biome-ignore lint/correctness/useExhaustiveDependencies: changing conversations intentionally retriggers the scroll reset.
React.useLayoutEffect(() => {
const vp = viewport();
if (vp) {
@@ -551,6 +576,7 @@ function ChatView({
}
}, [conv.id, viewport]);
// new activity → stick to bottom only if the user is already pinned there
// biome-ignore lint/correctness/useExhaustiveDependencies: message and running changes intentionally retrigger bottom anchoring.
React.useLayoutEffect(() => {
if (!atBottomRef.current) return;
const vp = viewport();
@@ -560,8 +586,12 @@ function ChatView({
// Per-conversation token total, live — same accounting as the main-agent
// console: completed runs' `result` sum + the in-progress run's latest `usage`.
const tokenTotal = React.useMemo(() => {
let i = 0, o = 0, cr = 0;
let li = 0, lo = 0, lcr = 0;
let i = 0,
o = 0,
cr = 0;
let li = 0,
lo = 0,
lcr = 0;
let turns = 0; // agent 循环轮次 = 模型调用次数（每次一条 kind='usage'）
for (const a of messages) {
if (a.kind === "result") {
@@ -576,7 +606,9 @@ function ChatView({
lcr = a.cache_read_tokens ?? 0;
}
}
const I = i + li, O = o + lo, CR = cr + lcr;
const I = i + li,
O = o + lo,
CR = cr + lcr;
return { i: I, o: O, cr: CR, turns, any: I + O + CR > 0 };
}, [messages]);
@@ -636,10 +668,10 @@ function ChatView({
return (
<>
{/* header: which agent + live + token meta */}
<div className="flex items-center gap-2 border-b px-4 py-2.5">
<div className="flex min-w-0 flex-wrap items-center gap-2 border-b px-4 py-2.5">
<Bot className="text-muted-foreground size-4 shrink-0" />
<span className="shrink-0 text-sm font-medium">{agent?.name ?? conv.agent_key}</span>
<span className="text-muted-foreground shrink-0 font-mono text-xs">{conv.agent_key}</span>
<span className="min-w-0 max-w-48 truncate text-sm font-medium">{agent?.name ?? conv.agent_key}</span>
<span className="text-muted-foreground hidden shrink-0 font-mono text-xs sm:inline">{conv.agent_key}</span>
{agent && !agent.builtin && (
<Badge variant="outline" className="shrink-0 px-1.5 py-0 text-[10px]">
自定义
@@ -649,14 +681,14 @@ function ChatView({
<span className="text-muted-foreground min-w-0 truncate text-xs">{agent.description}</span>
)}
{running && <LiveBadge />}
<div className="text-muted-foreground ml-auto flex shrink-0 items-center gap-3 text-xs">
<div className="text-muted-foreground ml-auto flex min-w-0 max-w-full items-center justify-end gap-x-3 gap-y-1 text-xs max-sm:w-full max-sm:flex-wrap">
{tokenTotal.turns > 0 && (
<span title="agent 循环轮次（模型调用次数）" className="tabular-nums">
{tokenTotal.turns} 轮
</span>
)}
{tokenTotal.any && (
<span title="input / cache(read) / output tokens" className="tabular-nums">
<span title="input / cache(read) / output tokens" className="min-w-0 truncate tabular-nums">
input {fmtTokens(tokenTotal.i)} · cache {fmtTokens(tokenTotal.cr)} · output {fmtTokens(tokenTotal.o)}
</span>
)}
@@ -664,10 +696,7 @@ function ChatView({
</div>
{/* messages */}
<ScrollArea
type="auto"
className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!"
>
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<div className="min-w-0 max-w-full px-4 py-3" ref={contentRef}>
{messages.length === 0 && !running ? (
<div className="text-muted-foreground py-10 text-center text-sm">
@@ -676,9 +705,7 @@ function ChatView({
) : (
<>
{hasMore && (
<div className="text-muted-foreground/70 pb-2 text-center text-[11px]">
向上滚动加载更早的消息…
</div>
<div className="text-muted-foreground/70 pb-2 text-center text-[11px]">向上滚动加载更早的消息…</div>
)}
<Transcript activity={messages} live={running} chat fetchDetail={fetchDetail} />
</>
@@ -711,8 +738,8 @@ function ChatView({
);
}
// ConversationItem is one row in the left rail — status-ish icon, title + agent
// subtitle, inline rename, hover delete. Modeled on the console's SessionItem.
// ConversationItem is one row in the left rail: title, agent subtitle, inline
// rename, pin marker, and a compact action menu.
function ConversationItem({
conv,
agent,
@@ -723,6 +750,8 @@ function ConversationItem({
onStartRename,
onRenameText,
onCommitRename,
onCancelRename,
onTogglePinned,
onDelete,
}: {
conv: Conversation;
@@ -734,8 +763,22 @@ function ConversationItem({
onStartRename: () => void;
onRenameText: (v: string) => void;
onCommitRename: () => void;
onCancelRename: () => void;
onTogglePinned: () => void;
onDelete: () => void;
}) {
const [deleteOpen, setDeleteOpen] = React.useState(false);
const renameInputRef = React.useRef<HTMLInputElement>(null);
const cancelRenameRef = React.useRef(false);
const pinned = conversationIsPinned(conv);
React.useEffect(() => {
if (!renaming) return;
cancelRenameRef.current = false;
renameInputRef.current?.focus();
renameInputRef.current?.select();
}, [renaming]);
return (
<div
className={cn(
@@ -745,13 +788,19 @@ function ConversationItem({
>
{renaming ? (
<input
autoFocus
ref={renameInputRef}
value={renameText}
onChange={(e) => onRenameText(e.target.value)}
onBlur={onCommitRename}
onBlur={() => {
if (!cancelRenameRef.current) onCommitRename();
}}
onKeyDown={(e) => {
if (e.key === "Enter") onCommitRename();
if (e.key === "Escape") onCommitRename();
if (e.key === "Escape") {
e.preventDefault();
cancelRenameRef.current = true;
onCancelRename();
}
}}
className="border-input bg-background min-w-0 flex-1 rounded-md border px-2 py-1 text-sm"
/>
@@ -763,26 +812,58 @@ function ConversationItem({
title="双击重命名"
className="min-w-0 flex-1 rounded-md px-2 py-1.5 text-left"
>
<div className="truncate text-sm">{conv.title || "新对话"}</div>
<div className="flex min-w-0 items-center gap-1.5">
{pinned && <PinIcon className="text-primary size-3 shrink-0" aria-label="已置顶" />}
<div className="truncate text-sm">{conv.title || "新对话"}</div>
</div>
<div className="text-muted-foreground flex min-w-0 items-center gap-1 text-[11px]">
<Bot className="size-3 shrink-0" />
<span className="min-w-0 truncate">{agent?.name ?? conv.agent_key}</span>
<span className="shrink-0">·</span>
<span className="shrink-0">{new Date(conv.created_at).toLocaleDateString("zh-CN", { month: "numeric", day: "numeric", hour: "2-digit", minute: "2-digit" })}</span>
<span className="shrink-0">
{new Date(conv.created_at).toLocaleDateString("zh-CN", {
month: "numeric",
day: "numeric",
hour: "2-digit",
minute: "2-digit",
})}
</span>
<span className="shrink-0 opacity-60">#{conv.id}</span>
</div>
</button>
)}
<AlertDialog>
<AlertDialogTrigger asChild>
<DropdownMenu>
<DropdownMenuTrigger asChild>
<Button
variant="ghost"
size="icon-sm"
className="text-muted-foreground hover:text-destructive shrink-0 opacity-0 transition-opacity group-hover:opacity-100"
className="text-muted-foreground shrink-0"
aria-label={`管理对话「${conv.title || "新对话"}」`}
>
<Trash2Icon className="size-3.5" />
<MoreHorizontalIcon />
</Button>
</AlertDialogTrigger>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuGroup>
<DropdownMenuItem onSelect={onStartRename}>
<PencilIcon />
重命名
</DropdownMenuItem>
<DropdownMenuItem onSelect={onTogglePinned}>
{pinned ? <PinOffIcon /> : <PinIcon />}
{pinned ? "取消置顶" : "置顶"}
</DropdownMenuItem>
</DropdownMenuGroup>
<DropdownMenuSeparator />
<DropdownMenuGroup>
<DropdownMenuItem variant="destructive" onSelect={() => setDeleteOpen(true)}>
<Trash2Icon />
删除
</DropdownMenuItem>
</DropdownMenuGroup>
</DropdownMenuContent>
</DropdownMenu>
<AlertDialog open={deleteOpen} onOpenChange={setDeleteOpen}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>删除对话「{conv.title || "新对话"}」？</AlertDialogTitle>
@@ -813,11 +894,24 @@ export default function ChatPage() {
>({});
const reloadConvs = React.useCallback(() => {
api.conversations().then(setConvs).catch(() => setConvs([]));
api
.conversations()
.then(setConvs)
.catch(() => setConvs([]));
}, []);
React.useEffect(() => {
api.agents().then(setAgents).catch(() => {});
api.llmProfiles().then(setProfiles).catch(() => {});
api
.agents()
.then(setAgents)
.catch(() => {
/* Agent metadata is optional for rendering an existing conversation. */
});
api
.llmProfiles()
.then(setProfiles)
.catch(() => {
/* The conversation remains usable without profile labels. */
});
reloadConvs();
}, [reloadConvs]);
@@ -856,6 +950,16 @@ export default function ChatPage() {
}
}
async function togglePinned(conversation: Conversation) {
const pinned = conversationIsPinned(conversation);
try {
await api.pinConversation(conversation.id, !pinned);
reloadConvs();
} catch (e) {
toast.error(`${pinned ? "取消置顶" : "置顶"}失败：${(e as Error).message}`);
}
}
function startRename(c: Conversation) {
setRenamingId(c.id);
setRenameText(c.title || "");
@@ -876,9 +980,9 @@ export default function ChatPage() {
return (
<div
data-content-padding="false"
className="flex h-[calc(100svh-3rem)] flex-col overflow-hidden p-4 md:h-[calc(100svh-4rem)] md:p-6"
className="flex h-[calc(100svh-3rem)] min-w-0 flex-col overflow-hidden p-3 sm:p-4 md:h-[calc(100svh-4rem)] md:p-6"
>
<div className="grid min-h-0 flex-1 grid-cols-[18rem_1fr] grid-rows-[minmax(0,1fr)] gap-4">
<div className="grid min-h-0 min-w-0 flex-1 grid-cols-1 grid-rows-[minmax(10rem,15rem)_minmax(0,1fr)] gap-3 md:grid-cols-[18rem_minmax(0,1fr)] md:grid-rows-[minmax(0,1fr)] md:gap-4">
{/* left: conversation list */}
<div className="bg-card flex flex-col overflow-hidden rounded-lg border">
<div className="border-b p-2">
@@ -888,9 +992,7 @@ export default function ChatPage() {
</div>
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<div className="flex min-w-0 flex-col gap-0.5 p-2">
{convs.length === 0 && (
<p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>
)}
{convs.length === 0 && <p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>}
{convs.map((c) => (
<ConversationItem
key={c.id}
@@ -903,6 +1005,8 @@ export default function ChatPage() {
onStartRename={() => startRename(c)}
onRenameText={setRenameText}
onCommitRename={commitRename}
onCancelRename={() => setRenamingId(null)}
onTogglePinned={() => void togglePinned(c)}
onDelete={() => del(c.id)}
/>
))}
@@ -931,7 +1035,12 @@ export default function ChatPage() {
// it on this render (switching to ChatView right away, before the
// async reloadConvs lands); reloadConvs then reconciles titles etc.
if (pending) setPendingByConv((p) => ({ ...p, [c.id]: pending }));
setConvs((prev) => (prev.some((x) => x.id === c.id) ? prev : [c, ...prev]));
setConvs((prev) => {
if (prev.some((item) => item.id === c.id)) return prev;
const firstUnpinned = prev.findIndex((item) => !conversationIsPinned(item));
const insertAt = firstUnpinned < 0 ? prev.length : firstUnpinned;
return [...prev.slice(0, insertAt), c, ...prev.slice(insertAt)];
});
setSelectedId(c.id);
reloadConvs();
}}
web/src/app/(main)/function/assets/page.tsx
+665 -255
View File
File diff suppressed because it is too large. Load diff
web/src/app/(main)/function/findings/page.tsx
+867 -539
View File
File diff suppressed because it is too large. Load diff
web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx
+729 -382
View File
File diff suppressed because it is too large. Load diff
web/src/app/(main)/function/tasks/page.tsx
+532 -326
View File
File diff suppressed because it is too large. Load diff
web/src/components/table-pagination.tsx
+21 -30
View File
@@ -1,22 +1,12 @@
"use client";
import * as React from "react";
import { useEffect, useRef } from "react";
import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
Pagination,
PaginationContent,
PaginationEllipsis,
PaginationItem,
} from "@/components/ui/pagination";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { Pagination, PaginationContent, PaginationEllipsis, PaginationItem } from "@/components/ui/pagination";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
interface TablePaginationProps {
page: number; // 1-based
@@ -49,19 +39,20 @@ export function TablePagination({
pageSizeOptions = [10, 20, 50],
}: TablePaginationProps) {
const totalPages = Math.max(1, Math.ceil(total / pageSize));
const from = total === 0 ? 0 : (page - 1) * pageSize + 1;
const to = Math.min(page * pageSize, total);
const safePage = Math.min(Math.max(1, page), totalPages);
const from = total === 0 ? 0 : (safePage - 1) * pageSize + 1;
const to = Math.min(safePage * pageSize, total);
const onPageChangeRef = useRef(onPageChange);
onPageChangeRef.current = onPageChange;
useEffect(() => {
if (page !== safePage) onPageChangeRef.current(safePage);
}, [page, safePage]);
return (
<div className="flex flex-wrap items-center justify-between gap-2 border-t px-4 py-2 lg:px-6">
<div className="text-muted-foreground flex items-center gap-2 text-xs">
<Select
value={String(pageSize)}
onValueChange={(v) => {
onPageSizeChange(Number(v));
onPageChange(1);
}}
>
<Select value={String(pageSize)} onValueChange={(v) => onPageSizeChange(Number(v))}>
<SelectTrigger size="sm" className="h-7 w-16">
<SelectValue />
</SelectTrigger>
@@ -90,14 +81,14 @@ export function TablePagination({
<Button
variant="ghost"
size="icon-sm"
disabled={page === 1}
onClick={() => onPageChange(page - 1)}
disabled={safePage === 1}
onClick={() => onPageChange(safePage - 1)}
aria-label="上一页"
>
<ChevronLeftIcon className="size-4" />
</Button>
</PaginationItem>
{pageWindows(page, totalPages).map((p, i) =>
{pageWindows(safePage, totalPages).map((p, i) =>
p === "..." ? (
<PaginationItem key={`el-${i}`}>
<PaginationEllipsis />
@@ -105,10 +96,10 @@ export function TablePagination({
) : (
<PaginationItem key={p}>
<Button
variant={p === page ? "outline" : "ghost"}
variant={p === safePage ? "outline" : "ghost"}
size="icon-sm"
onClick={() => onPageChange(p)}
aria-current={p === page ? "page" : undefined}
aria-current={p === safePage ? "page" : undefined}
>
{p}
</Button>
@@ -119,8 +110,8 @@ export function TablePagination({
<Button
variant="ghost"
size="icon-sm"
disabled={page === totalPages}
onClick={() => onPageChange(page + 1)}
disabled={safePage === totalPages}
onClick={() => onPageChange(safePage + 1)}
aria-label="下一页"
>
<ChevronRightIcon className="size-4" />
web/src/components/task-template-controls.tsx
+455
View File
@@ -0,0 +1,455 @@
"use client";
import * as React from "react";
import { LibraryIcon, PlusIcon, SaveIcon, Settings2Icon, Trash2Icon } from "lucide-react";
import { toast } from "sonner";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Spinner } from "@/components/ui/spinner";
import { Textarea } from "@/components/ui/textarea";
import { api } from "@/lib/api";
import type { TaskTemplate } from "@/lib/types";
import { cn } from "@/lib/utils";
interface TemplateDraft {
name: string;
description: string;
goal: string;
}
interface TaskTemplateManagerProps {
open: boolean;
onOpenChange: (open: boolean) => void;
templates: TaskTemplate[];
seed: Pick<TemplateDraft, "description" | "goal"> | null;
onCreated: (template: TaskTemplate) => void;
onUpdated: (template: TaskTemplate) => void;
onDeleted: (id: number) => void;
}
const emptyDraft = (): TemplateDraft => ({ name: "", description: "", goal: "" });
function templateDraft(template: TaskTemplate): TemplateDraft {
return {
name: template.name,
description: template.description,
goal: template.goal,
};
}
function TaskTemplateManager({
open,
onOpenChange,
templates,
seed,
onCreated,
onUpdated,
onDeleted,
}: TaskTemplateManagerProps) {
const [selectedID, setSelectedID] = React.useState<number | null>(null);
const [draft, setDraft] = React.useState<TemplateDraft>(emptyDraft);
const [saving, setSaving] = React.useState(false);
const [deleteOpen, setDeleteOpen] = React.useState(false);
const [deleting, setDeleting] = React.useState(false);
const wasOpen = React.useRef(false);
React.useEffect(() => {
if (open && !wasOpen.current) {
if (seed) {
setSelectedID(null);
setDraft({ name: "", description: seed.description, goal: seed.goal });
} else if (templates[0]) {
setSelectedID(templates[0].id);
setDraft(templateDraft(templates[0]));
} else {
setSelectedID(null);
setDraft(emptyDraft());
}
}
wasOpen.current = open;
}, [open, seed, templates]);
const selectTemplate = (template: TaskTemplate) => {
setSelectedID(template.id);
setDraft(templateDraft(template));
};
const startNew = () => {
setSelectedID(null);
setDraft(emptyDraft());
};
const updateDraft = (field: keyof TemplateDraft, value: string) => {
setDraft((current) => ({ ...current, [field]: value }));
};
async function save() {
const input = {
name: draft.name.trim(),
description: draft.description.trim(),
goal: draft.goal.trim(),
};
if (!input.name || !input.description || !input.goal) {
toast.error("请填写模板名称、描述和目标");
return;
}
setSaving(true);
try {
if (selectedID == null) {
const created = await api.createTaskTemplate(input);
onCreated(created);
setSelectedID(created.id);
setDraft(templateDraft(created));
toast.success("模板已创建");
} else {
const updated = await api.updateTaskTemplate(selectedID, input);
onUpdated(updated);
setDraft(templateDraft(updated));
toast.success("模板已更新");
}
} catch (error) {
toast.error(`保存失败：${(error as Error).message}`);
} finally {
setSaving(false);
}
}
async function remove() {
if (selectedID == null) return;
const deletedID = selectedID;
setDeleting(true);
try {
await api.deleteTaskTemplate(deletedID);
onDeleted(deletedID);
const next = templates.find((template) => template.id !== deletedID);
if (next) {
setSelectedID(next.id);
setDraft(templateDraft(next));
} else {
startNew();
}
setDeleteOpen(false);
toast.success("模板已删除");
} catch (error) {
toast.error(`删除失败：${(error as Error).message}`);
} finally {
setDeleting(false);
}
}
let saveLabel = saving ? "保存中" : "保存修改";
if (!saving && selectedID == null) saveLabel = "创建模板";
return (
<>
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="grid max-h-[90svh] max-w-4xl grid-rows-[auto_minmax(0,1fr)_auto] gap-0 overflow-hidden p-0">
<DialogHeader className="border-b px-6 py-5">
<DialogTitle>任务模板管理</DialogTitle>
<DialogDescription>模板只保存描述与目标；修改不会影响已经创建的任务。</DialogDescription>
</DialogHeader>
<div className="grid min-h-0 overflow-y-auto md:grid-cols-[15rem_minmax(0,1fr)] md:overflow-hidden">
<div className="flex min-h-0 flex-col border-b p-3 md:border-r md:border-b-0">
<Button type="button" variant="outline" className="w-full" onClick={startNew}>
<PlusIcon data-icon="inline-start" />
新建模板
</Button>
<ScrollArea className="mt-2 max-h-44 md:max-h-[27rem]">
<div className="flex flex-col gap-1 pr-2">
{templates.length === 0 && (
<p className="px-2 py-6 text-center text-muted-foreground text-sm">暂无模板</p>
)}
{templates.map((template) => (
<button
key={template.id}
type="button"
className={cn(
"min-w-0 rounded-md px-2.5 py-2 text-left transition-colors",
selectedID === template.id ? "bg-accent text-accent-foreground" : "hover:bg-accent/50",
)}
onClick={() => selectTemplate(template)}
>
<span className="block truncate font-medium text-sm">{template.name}</span>
<span className="block truncate text-muted-foreground text-xs">{template.description}</span>
</button>
))}
</div>
</ScrollArea>
</div>
<ScrollArea className="max-h-[58svh]">
<FieldGroup className="p-6">
<Field>
<FieldLabel htmlFor="task-template-name">模板名称</FieldLabel>
<Input
id="task-template-name"
value={draft.name}
maxLength={120}
placeholder="例如：外部 Web 渗透"
onChange={(event) => updateDraft("name", event.target.value)}
/>
</Field>
<Field>
<FieldLabel htmlFor="task-template-description">描述</FieldLabel>
<Textarea
id="task-template-description"
className="min-h-28"
value={draft.description}
placeholder="测试对象与背景"
onChange={(event) => updateDraft("description", event.target.value)}
/>
</Field>
<Field>
<FieldLabel htmlFor="task-template-goal">目标</FieldLabel>
<Textarea
id="task-template-goal"
className="min-h-28"
value={draft.goal}
placeholder="任务需要达成的目标"
onChange={(event) => updateDraft("goal", event.target.value)}
/>
</Field>
</FieldGroup>
</ScrollArea>
</div>
<DialogFooter className="border-t px-6 py-4">
{selectedID != null && (
<Button type="button" variant="destructive" className="sm:mr-auto" onClick={() => setDeleteOpen(true)}>
<Trash2Icon data-icon="inline-start" />
删除模板
</Button>
)}
<Button type="button" variant="outline" onClick={() => onOpenChange(false)}>
关闭
</Button>
<Button type="button" disabled={saving} onClick={() => void save()}>
{saving ? <Spinner data-icon="inline-start" /> : <SaveIcon data-icon="inline-start" />}
{saveLabel}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
<AlertDialog open={deleteOpen} onOpenChange={setDeleteOpen}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>删除模板「{draft.name || "未命名模板"}」？</AlertDialogTitle>
<AlertDialogDescription>已由该模板创建的任务不会受到影响。</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={deleting}>取消</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
disabled={deleting}
onClick={(event) => {
event.preventDefault();
void remove();
}}
>
{deleting && <Spinner data-icon="inline-start" />}
{deleting ? "删除中" : "删除"}
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</>
);
}
interface TaskTemplateControlsProps {
description: string;
goal: string;
selectedTemplateID: number | null;
onSelectedTemplateIDChange: (id: number | null) => void;
onApply: (template: TaskTemplate) => void;
portalContainer?: React.RefObject<HTMLElement | null>;
}
export function TaskTemplateControls({
description,
goal,
selectedTemplateID,
onSelectedTemplateIDChange,
onApply,
portalContainer,
}: TaskTemplateControlsProps) {
const [templates, setTemplates] = React.useState<TaskTemplate[]>([]);
const [loading, setLoading] = React.useState(true);
const [pendingTemplate, setPendingTemplate] = React.useState<TaskTemplate | null>(null);
const [managerOpen, setManagerOpen] = React.useState(false);
const [managerSeed, setManagerSeed] = React.useState<Pick<TemplateDraft, "description" | "goal"> | null>(null);
const loadTemplates = React.useCallback(async () => {
setLoading(true);
try {
setTemplates(await api.taskTemplates());
} catch (error) {
setTemplates([]);
toast.error(`加载模板失败：${(error as Error).message}`);
} finally {
setLoading(false);
}
}, []);
React.useEffect(() => {
void loadTemplates();
}, [loadTemplates]);
const templateIDs = React.useMemo(() => templates.map((template) => String(template.id)), [templates]);
const templatesByID = React.useMemo(
() => new Map(templates.map((template) => [String(template.id), template])),
[templates],
);
const applyTemplate = (template: TaskTemplate) => {
onApply(template);
onSelectedTemplateIDChange(template.id);
setPendingTemplate(null);
};
const chooseTemplate = (value: string | null) => {
if (!value) {
onSelectedTemplateIDChange(null);
return;
}
const template = templatesByID.get(value);
if (!template) return;
const hasContent = description.trim() !== "" || goal.trim() !== "";
const changesContent = description !== template.description || goal !== template.goal;
if (hasContent && changesContent) {
setPendingTemplate(template);
return;
}
applyTemplate(template);
};
const openManager = (seed: Pick<TemplateDraft, "description" | "goal"> | null) => {
setManagerSeed(seed);
setManagerOpen(true);
};
const upsertTemplate = (template: TaskTemplate) => {
setTemplates((current) => [template, ...current.filter((item) => item.id !== template.id)]);
};
const deleteTemplate = (id: number) => {
setTemplates((current) => current.filter((template) => template.id !== id));
if (selectedTemplateID === id) onSelectedTemplateIDChange(null);
};
let pickerPlaceholder = loading ? "正在加载模板" : "暂无任务模板";
if (!loading && templates.length > 0) pickerPlaceholder = "搜索并选择任务模板";
return (
<>
<Field>
<div className="flex flex-wrap items-center justify-between gap-2">
<FieldLabel htmlFor="task-template-picker">任务模板</FieldLabel>
<div className="flex flex-wrap gap-2">
<Button type="button" variant="ghost" size="sm" onClick={() => openManager(null)}>
<Settings2Icon data-icon="inline-start" />
管理模板
</Button>
<Button
type="button"
variant="ghost"
size="sm"
disabled={!description.trim() || !goal.trim()}
onClick={() => openManager({ description, goal })}
>
<SaveIcon data-icon="inline-start" />
另存为模板
</Button>
</div>
</div>
<Combobox
items={templateIDs}
itemToStringValue={(id) => templatesByID.get(id)?.name ?? id}
value={selectedTemplateID == null ? null : String(selectedTemplateID)}
onValueChange={chooseTemplate}
>
<ComboboxInput
id="task-template-picker"
className="w-full"
placeholder={pickerPlaceholder}
disabled={loading || templates.length === 0}
showClear
/>
<ComboboxContent portalContainer={portalContainer}>
<ComboboxEmpty>没有匹配的模板</ComboboxEmpty>
<ComboboxList>
{(id) => {
const template = templatesByID.get(id);
return (
<ComboboxItem key={id} value={id}>
<LibraryIcon />
<div className="min-w-0 flex-1">
<p className="truncate font-medium text-sm">{template?.name ?? `模板 #${id}`}</p>
{template?.description && (
<p className="truncate text-muted-foreground text-xs">{template.description}</p>
)}
</div>
</ComboboxItem>
);
}}
</ComboboxList>
</ComboboxContent>
</Combobox>
<FieldDescription>选择后会复制模板的描述和目标，不与模板保持关联。</FieldDescription>
</Field>
<AlertDialog open={pendingTemplate != null} onOpenChange={(open) => !open && setPendingTemplate(null)}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>使用模板「{pendingTemplate?.name}」？</AlertDialogTitle>
<AlertDialogDescription>当前已填写的描述和目标将被模板内容覆盖。</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>取消</AlertDialogCancel>
<AlertDialogAction onClick={() => pendingTemplate && applyTemplate(pendingTemplate)}>
覆盖并使用
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
<TaskTemplateManager
open={managerOpen}
onOpenChange={setManagerOpen}
templates={templates}
seed={managerSeed}
onCreated={upsertTemplate}
onUpdated={upsertTemplate}
onDeleted={deleteTemplate}
/>
</>
);
}
web/src/components/ui/alert.tsx
+3 -3
View File
@@ -4,7 +4,7 @@ import { cva, type VariantProps } from "class-variance-authority"
import { cn } from "@/lib/utils"
const alertVariants = cva(
"group/alert relative grid w-full gap-0.5 rounded-lg border px-2.5 py-2 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4",
"group/alert relative grid w-full gap-0.5 rounded-lg border px-2.5 py-2 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_minmax(0,1fr)] has-[>svg]:gap-x-2 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4",
{
variants: {
variant: {
@@ -39,7 +39,7 @@ function AlertTitle({ className, ...props }: React.ComponentProps<"div">) {
<div
data-slot="alert-title"
className={cn(
"font-medium group-has-[>svg]/alert:col-start-2 [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground",
"min-w-0 font-medium group-has-[>svg]/alert:col-start-2 [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground",
className
)}
{...props}
@@ -55,7 +55,7 @@ function AlertDescription({
<div
data-slot="alert-description"
className={cn(
"text-sm text-balance text-muted-foreground md:text-pretty [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground [&_p:not(:last-child)]:mb-4",
"min-w-0 text-sm text-balance text-muted-foreground [overflow-wrap:anywhere] md:text-pretty [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground [&_p:not(:last-child)]:mb-4",
className
)}
{...props}
web/src/lib/api.ts
+51 -26
View File
@@ -13,9 +13,11 @@ import type {
Asset,
AssetNode,
Audit,
BatchControlItem,
ChatAttachment,
CommandRecord,
Company,
CompanyScopeRule,
Conversation,
ConvTokenSummary,
CoverageAssetRefs,
@@ -25,6 +27,8 @@ import type {
DeleteTaskResult,
Edge,
Finding,
FindingDeepenResponse,
FindingGroupsPage,
FindingQuery,
FindingStats,
FindingStatus,
@@ -39,21 +43,24 @@ import type {
LLMTask,
MCPServer,
MCPTool,
MissingSkill,
ModelTokenStat,
PromptVar,
PromptVersion,
SessionTokenUsage,
Settings,
Severity,
SkillCall,
SkillItem,
MissingSkill,
SSProject,
SSTask,
Stats,
Task,
TaskAssetView,
TaskLLMResolutions,
TaskNode,
TaskScopeRow,
TaskTemplate,
TokenTotal,
TokenUsage,
Tool,
@@ -179,6 +186,7 @@ export const api = {
goal: string;
llmProfileIds?: number[];
sourceTaskIds?: string[];
companyIds?: number[];
timeoutSeconds?: number;
seedFirstIntent?: boolean;
planHeartbeatSeconds?: number;
@@ -188,10 +196,17 @@ export const api = {
goal: input.goal,
llm_profile_ids: input.llmProfileIds ?? [],
source_task_ids: input.sourceTaskIds ?? [],
company_ids: input.companyIds ?? [],
timeout_seconds: input.timeoutSeconds ?? 0,
seed_first_intent: input.seedFirstIntent ?? false,
plan_heartbeat_seconds: input.planHeartbeatSeconds ?? 0, // 0 = 后端归一到默认 600(10min)
}),
taskTemplates: () => get<{ templates: TaskTemplate[] }>("/task-templates").then((r) => arr(r.templates)),
createTaskTemplate: (input: Pick<TaskTemplate, "name" | "description" | "goal">) =>
post<TaskTemplate>("/task-templates", input),
updateTaskTemplate: (id: number, input: Partial<Pick<TaskTemplate, "name" | "description" | "goal">>) =>
patch<TaskTemplate>(`/task-templates/${id}`, input),
deleteTaskTemplate: (id: number) => del<{ deleted: number }>(`/task-templates/${id}`),
updateTaskLLMProfiles: (id: string, llmProfileIds: number[], activeLLMProfileId?: number) =>
put<{
id: string;
@@ -206,13 +221,21 @@ export const api = {
}),
deleteTask: (id: string, options: DeleteTaskOptions) => del<DeleteTaskResult>(`/tasks/${id}`, options),
controlTask: (id: string, action: "pause" | "resume") =>
post<{ id: string; paused: boolean }>(`/tasks/${id}/control`, { action }),
post<{ id: string; paused: boolean; queued: boolean; status: string }>(`/tasks/${id}/control`, { action }),
controlTasksBatch: (taskIds: string[], action: "pause" | "resume") =>
post<{ items: BatchControlItem[] }>("/tasks/control/batch", { task_ids: taskIds, action }),
controlIntent: (taskId: string, intentId: string, action: "pause" | "resume" | "cancel") =>
post<{
id: number;
state: "paused" | "open" | "cancelled";
deleted?: { intents: number; facts: number; findings: number; activities: number };
}>(`/tasks/${taskId}/intents/${intentId}/control`, { action }),
controlIntentsBatch: (taskId: string, intentIds: string[], action: "pause" | "resume") =>
post<{ items: BatchControlItem[] }>(`/tasks/${taskId}/intents/control/batch`, {
intent_ids: intentIds,
action,
}),
taskLLMResolution: (id: string) => get<TaskLLMResolutions>(`/tasks/${id}/llm/resolution`),
// 重跑一条没跑成功的意图(blocked/exhausted/stopped)：置回 open，worker 会重新认领、从头再跑。
rerunIntent: (taskId: string, intentId: string) =>
post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/${intentId}/rerun`),
@@ -308,38 +331,22 @@ export const api = {
// ---- companies (企业 + 资产范围；归属唯一来源) ----
companies: () => get<Company[]>("/companies").then(arr),
createCompany: (name: string, logo: string, scope: string) =>
createCompany: (name: string, scope: CompanyScopeRule[]) =>
post<{ id: number; created: boolean; scope_added?: number; scope_invalid?: number; scope_errors?: string[] }>(
"/companies",
{
name,
logo,
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
},
),
addCompanyScope: (id: number, scope: string, reason = "") =>
addCompanyScope: (id: number, scope: CompanyScopeRule[], reason = "") =>
post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, {
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
reason,
}),
updateCompanyScope: (id: number, scope: string, reason = "") =>
updateCompanyScope: (id: number, scope: CompanyScopeRule[], reason = "") =>
post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, {
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
reason,
reset: true,
}),
@@ -361,6 +368,15 @@ export const api = {
if (q.sort) p.set("sort", q.sort);
return get<FindingsPage>(`/exploration/findings?${p.toString()}`);
},
findingGroups: (q: FindingQuery) => {
const p = new URLSearchParams({ page: String(q.page), limit: String(q.pageSize) });
if (q.severity && q.severity !== "all") p.set("severity", q.severity);
if (q.status && q.status !== "all") p.set("status", q.status);
if (q.vulnclass && q.vulnclass !== "all") p.set("vulnclass", q.vulnclass);
if (q.task && q.task !== "all") p.set("task_id", q.task);
if (q.sort) p.set("sort", q.sort);
return get<FindingGroupsPage>(`/exploration/findings/groups?${p.toString()}`);
},
findingStats: () => get<FindingStats>("/exploration/findings/stats"),
// exportFindings 触发发现页导出并下载文件。scope=selected 时传 ids(finding_id 列表);
// scope=filtered 时传当前筛选(沿用 FindingQuery 的筛选字段);scope=all 忽略筛选。
@@ -415,10 +431,15 @@ export const api = {
) => patch<Finding>(`/exploration/findings/${id}`, fields),
// 删除漏洞:移除 findings 记录 + 来源探索节点(从发现列表/任务发现 Tab/探索图一并消失)。
deleteFinding: (id: string) => del<{ deleted: boolean; id: number }>(`/exploration/findings/${id}`),
deepenFinding: (id: string, description: string) =>
post<FindingDeepenResponse>(`/exploration/findings/${id}/deepen`, { description }),
intents: (task?: string) => get<TaskNode[]>(`/exploration/intents${tq(task)}`).then(arr),
tokenStats: (task?: string) =>
get<{ workers: TokenUsage[]; total: TokenTotal }>(`/exploration/tokens${tq(task)}`).then((r) => ({
get<{ workers: TokenUsage[]; sessions: SessionTokenUsage[]; total: TokenTotal }>(
`/exploration/tokens${tq(task)}`,
).then((r) => ({
workers: arr(r.workers),
sessions: arr(r.sessions),
total: r.total,
})),
tokenDaily: (days = 30) => get<DailyTokenBucket[]>(`/tokens/daily?days=${days}`).then(arr),
@@ -505,6 +526,7 @@ export const api = {
},
chat: (message: string, task?: string, attachments?: ChatAttachment[]) =>
post<{ reply: string; mode: string }>(`/chat${tq(task)}`, { message, attachments }),
chatStatus: (taskId: string) => get<{ running: boolean }>(`/tasks/${taskId}/chat/status`),
// 方式1 文件上传:落到会话/任务工作目录 uploads/，返回可供 agent Read 的相对路径。
chatUpload: async (scope: "task" | "session" | "staging", id: string, files: File[]) => {
if (MOCK)
@@ -619,7 +641,10 @@ export const api = {
conversations: () => get<{ conversations: Conversation[] }>("/conversations").then((r) => arr(r.conversations)),
createConversation: (agent_key: string, title = "", llm_profile_id?: number | null) =>
post<Conversation>("/conversations", { agent_key, title, llm_profile_id: llm_profile_id ?? null }),
renameConversation: (id: number, title: string) => patch<{ ok: boolean }>(`/conversations/${id}`, { title }),
updateConversation: (id: number, input: { title?: string; pinned?: boolean }) =>
patch<Conversation>(`/conversations/${id}`, input),
renameConversation: (id: number, title: string) => patch<Conversation>(`/conversations/${id}`, { title }),
pinConversation: (id: number, pinned: boolean) => patch<Conversation>(`/conversations/${id}`, { pinned }),
updateConversationProfile: (id: number, llm_profile_id: number | null) =>
patch<{ ok: boolean }>(`/conversations/${id}/profile`, { llm_profile_id }),
deleteConversation: (id: number) => del<{ deleted: number }>(`/conversations/${id}`),
web/src/lib/mock/data.ts
+49 -3
View File
@@ -7,8 +7,8 @@ import type {
Activity, Agent, AgentDetail, Asset, Audit, Company, Conversation,
ConvTokenSummary, DailyTokenBucket, Edge, Finding, InterceptApprovalRow,
InterceptPending, InterceptRule, LLMPoolStatus, LLMProfile, MCPServer, MCPTool, PromptVar,
PromptVersion, Settings, SkillCall, SkillItem, MissingSkill, Stats, TaskNode, Task, TokenTotal,
TokenUsage, Tool, TrafficResp, TrafficDetail, TrafficHost, LLMTask,
PromptVersion, Settings, SessionTokenUsage, SkillCall, SkillItem, MissingSkill, Stats, TaskNode, Task, TokenTotal,
TaskTemplate, TokenUsage, Tool, TrafficResp, TrafficDetail, TrafficHost, LLMTask,
} from "@/lib/types";
const T = (iso: string) => iso; // readability helper for timestamps
@@ -93,6 +93,25 @@ export const tasks: Task[] = [
},
];
export const taskTemplates: TaskTemplate[] = [
{
id: 1,
name: "外部 Web 渗透",
description: "对目标互联网暴露面开展黑盒渗透测试，覆盖站点、接口和常见管理入口。",
goal: "识别并验证可造成未授权访问、敏感数据泄露或服务器失陷的安全问题。",
created_at: T("2026-07-20T08:00:00Z"),
updated_at: T("2026-07-25T08:00:00Z"),
},
{
id: 2,
name: "API 越权专项",
description: "围绕目标 API 的身份认证、对象级授权和角色边界开展专项测试。",
goal: "确认是否存在 IDOR、水平越权、垂直越权及批量数据访问风险。",
created_at: T("2026-07-19T08:00:00Z"),
updated_at: T("2026-07-24T08:00:00Z"),
},
];
export const getTask = (id: string) => tasks.find((t) => t.id === id);
const grandTotal: TokenTotal = tasks.reduce(
@@ -516,6 +535,33 @@ export const tokenWorkers: TokenUsage[] = [
{ worker: "work#13", input_tokens: 144000, output_tokens: 9200, cache_read_tokens: 96000, cache_write_tokens: 14000 },
];
// Mirrors TokenStatsBySession: completed usage is keyed by the stable UI session
// (main | plan | intent:<id>) instead of the reusable work#N executor name.
export const tokenSessions: SessionTokenUsage[] = (() => {
const totals = new Map<string, SessionTokenUsage>();
for (const item of activity) {
if (item.kind !== "result") continue;
let session = "";
if (item.worker === "mainagent") session = "main";
else if (item.worker === "planner") session = "plan";
else if (item.intent_id) session = `intent:${item.intent_id}`;
if (!session) continue;
const current = totals.get(session) ?? {
session,
input_tokens: 0,
output_tokens: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
};
current.input_tokens += item.input_tokens ?? 0;
current.output_tokens += item.output_tokens ?? 0;
current.cache_read_tokens += item.cache_read_tokens ?? 0;
current.cache_write_tokens += item.cache_write_tokens ?? 0;
totals.set(session, current);
}
return [...totals.values()];
})();
export const tokenTotal: TokenTotal = grandTotal;
export const dailyTokens: DailyTokenBucket[] = Array.from({ length: 30 }, (_, i) => {
@@ -755,7 +801,7 @@ export const interceptHistory: InterceptApprovalRow[] = [
// ── Conversations (chat) ─────────────────────────────────────────────────────
export const conversations: Conversation[] = [
{ id: 1, agent_key: "mainagent", title: "acme 后台切入点讨论", llm_profile_id: 1, created_at: T("2026-07-25T10:00:00Z"), updated_at: T("2026-07-26T03:40:00Z") },
{ id: 1, agent_key: "mainagent", title: "acme 后台切入点讨论", llm_profile_id: 1, pinned: true, pinned_at: T("2026-07-26T04:00:00Z"), created_at: T("2026-07-25T10:00:00Z"), updated_at: T("2026-07-26T03:40:00Z") },
{ id: 2, agent_key: "mainagent", title: "API 越权面梳理", llm_profile_id: 2, created_at: T("2026-07-24T16:00:00Z"), updated_at: T("2026-07-24T17:20:00Z") },
{ id: 3, agent_key: "mainagent", title: "内网横向 · DMZ→域控攻击路径规划", llm_profile_id: 1, created_at: T("2026-07-25T15:45:00Z"), updated_at: T("2026-07-26T00:25:00Z") },
];
web/src/lib/mock/handler.ts
+429 -22
View File
@@ -2,7 +2,17 @@
// 未命中的一律返回安全默认（[] / {} / {ok:true}），保证任何页面都不崩。
// 只在 NEXT_PUBLIC_MOCK=1 时经由 api.ts 的 http() 短路进入这里。
import type { Task } from "../types";
import type {
Asset,
BatchControlItem,
Company,
CompanyScopeRule,
Conversation,
ScopeRow,
Task,
TaskLLMResolution,
TaskTemplate,
} from "../types";
import * as D from "./data";
const delay = (ms = 120) => new Promise((r) => setTimeout(r, ms));
@@ -12,8 +22,148 @@ const delay = (ms = 120) => new Promise((r) => setTimeout(r, ms));
const mockTasks = structuredClone(D.tasks);
const mockFindings = structuredClone(D.findings);
const mockLLMRecords = structuredClone(D.llmRecords);
const mockTaskTemplates = structuredClone(D.taskTemplates);
const mockConversations = structuredClone(D.conversations);
const mockIntents = structuredClone(D.intents);
const mockCompanies = structuredClone(D.companies);
const mockAssets = structuredClone(D.assets);
let mockActiveTask = D.ACTIVE_TASK;
function mockAssetCounts(): Record<string, number> {
return mockAssets.reduce<Record<string, number>>((counts, asset) => {
counts[asset.type] = (counts[asset.type] ?? 0) + 1;
return counts;
}, {});
}
function mockScopeRows(companyID: number, input: unknown): ScopeRow[] {
if (!Array.isArray(input)) return [];
let nextID =
mockCompanies.flatMap((company) => company.scope ?? []).reduce((max, row) => Math.max(max, row.id), 0) + 1;
const rows: ScopeRow[] = [];
for (const candidate of input) {
if (!candidate || typeof candidate !== "object") continue;
const rule = candidate as Partial<CompanyScopeRule>;
const kind = rule.kind;
const value = String(rule.value ?? "").trim();
if (!kind || !value) continue;
const row: ScopeRow = { id: nextID++, company_id: companyID, kind, raw: value };
if (kind === "domain") row.domain = value.toLowerCase();
else if (kind === "ip" || kind === "cidr") row.net = value;
else row.value = value;
rows.push(row);
}
return rows;
}
function mockProfileResolution(profileID: number | undefined, source: TaskLLMResolution["source"]): TaskLLMResolution {
const profile = D.llmProfiles.find((item) => Number(item.id) === profileID);
if (!profile) {
return { name: "", format: "", model: "", source, available: false, reason: "LLM 配置不存在" };
}
return {
profile_id: Number(profile.id),
name: profile.name,
format: profile.format,
model: profile.model,
source,
available: Boolean(profile.api_key_hint),
reason: profile.api_key_hint ? undefined : "LLM 配置未设置 API Key",
};
}
function mockRoleResolution(task: Task, agentKey: "mainagent" | "planner" | "worker"): TaskLLMResolution {
if (task.llm_profile_ids?.length) {
if (task.active_llm_profile_id === undefined) {
return {
name: "",
format: "",
model: "",
source: "task_chain",
available: false,
reason: "任务 LLM 配置链额度已耗尽",
};
}
return mockProfileResolution(task.active_llm_profile_id, "task_chain");
}
const agent = D.agents.find((item) => item.key === agentKey);
if (agent?.llm_profile_id) {
const bound = mockProfileResolution(agent.llm_profile_id, "agent_binding");
if (bound.available) return bound;
}
const globalProfile = D.llmProfiles.find((item) => item.is_default);
if (globalProfile) return mockProfileResolution(Number(globalProfile.id), "global_profile");
return {
name: "全局配置",
format: D.llmConfig.provider,
model: D.llmConfig.model,
source: "environment",
available: true,
};
}
function bodyIDs(value: unknown): string[] {
if (!Array.isArray(value)) return [];
const ids = [...new Set(value.map((item) => String(item)).filter(Boolean))];
if (ids.length > 100) throw new Error("批量操作最多支持 100 个 ID");
return ids;
}
function controlMockIntent(id: string, action: "pause" | "resume"): BatchControlItem {
const intent = mockIntents.find((item) => item.id === id);
if (!intent) return { id, ok: false, error: "意图不存在" };
const requiredState = action === "pause" ? "running" : "paused";
if (intent.inherited || intent.state !== requiredState) {
return { id, ok: false, state: intent.state, error: "Worker 状态已变化" };
}
intent.state = action === "pause" ? "paused" : "open";
return { id, ok: true, state: intent.state };
}
function controlMockTask(id: string, action: "pause" | "resume"): BatchControlItem {
const task = mockTasks.find((item) => item.id === id);
if (!task) return { id, ok: false, error: "任务不存在" };
if (task.status === "done" || task.status === "failed" || task.status === "timeout") {
return { id, ok: false, status: task.status, error: "终态任务不可控制" };
}
if (action === "pause") {
if (task.paused || task.status === "paused") {
return { id, ok: false, status: task.status, error: "任务已经暂停" };
}
task.paused = true;
task.queued = false;
task.status = "paused";
task.engine_mode = "paused";
} else {
if (!task.paused && task.status !== "paused") {
return { id, ok: false, status: task.status, error: "任务未暂停" };
}
task.paused = false;
task.queued = false;
task.status = "running";
task.engine_mode = "exploring";
}
return { id, ok: true, status: task.status, queued: false };
}
function sortMockConversations() {
mockConversations.sort((a, b) => {
const aPinned = a.pinned_at ? 1 : 0;
const bPinned = b.pinned_at ? 1 : 0;
if (aPinned !== bPinned) return bPinned - aPinned;
const aTime = a.pinned_at ?? a.updated_at;
const bTime = b.pinned_at ?? b.updated_at;
return bTime.localeCompare(aTime) || b.id - a.id;
});
}
function normalizedTemplateName(value: unknown): string {
return String(value ?? "")
.trim()
.split(/\s+/)
.join(" ");
}
function parseBody(body?: BodyInit | null): Record<string, unknown> {
if (typeof body !== "string") return {};
try {
@@ -50,6 +200,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
const now = new Date();
const profileIDs = [...((b.llm_profile_ids as number[] | undefined) ?? [])];
const sourceTaskIDs = [...((b.source_task_ids as string[] | undefined) ?? [])];
const companyIDs = [...((b.company_ids as number[] | undefined) ?? [])];
const created: Task = {
id,
description: String(b.description ?? "新任务"),
@@ -70,12 +221,64 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
active_llm_profile_id: profileIDs[0],
llm_failover_state: profileIDs.length ? "ready" : "default",
source_task_ids: sourceTaskIDs,
company_ids: companyIDs,
};
for (const item of mockTasks) item.active = false;
mockTasks.unshift(created);
mockActiveTask = id;
return created;
}
if (path === "/task-templates" && m === "GET") return { templates: mockTaskTemplates };
if (path === "/task-templates" && m === "POST") {
const now = new Date().toISOString();
const name = normalizedTemplateName(b.name);
const description = String(b.description ?? "").trim();
const goal = String(b.goal ?? "").trim();
if (!name || !description || !goal) throw new Error("请填写模板名称、描述和目标");
if (
mockTaskTemplates.some((template) => normalizedTemplateName(template.name).toLowerCase() === name.toLowerCase())
) {
throw new Error("模板名称已存在");
}
const nextID = mockTaskTemplates.reduce((max, template) => Math.max(max, template.id), 0) + 1;
const created: TaskTemplate = {
id: nextID,
name,
description,
goal,
created_at: now,
updated_at: now,
};
mockTaskTemplates.unshift(created);
return created;
}
if (seg[0] === "task-templates" && seg.length === 2 && m === "PATCH") {
const template = mockTaskTemplates.find((item) => item.id === Number(seg[1]));
if (!template) return {};
const name = typeof b.name === "string" ? normalizedTemplateName(b.name) : template.name;
const description = typeof b.description === "string" ? b.description.trim() : template.description;
const goal = typeof b.goal === "string" ? b.goal.trim() : template.goal;
if (!name || !description || !goal) throw new Error("请填写模板名称、描述和目标");
if (
mockTaskTemplates.some(
(item) => item.id !== template.id && normalizedTemplateName(item.name).toLowerCase() === name.toLowerCase(),
)
) {
throw new Error("模板名称已存在");
}
template.name = name;
template.description = description;
template.goal = goal;
template.updated_at = new Date().toISOString();
mockTaskTemplates.sort((a, b) => b.updated_at.localeCompare(a.updated_at) || b.id - a.id);
return template;
}
if (seg[0] === "task-templates" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockTaskTemplates.findIndex((item) => item.id === id);
if (index >= 0) mockTaskTemplates.splice(index, 1);
return { deleted: id };
}
if (seg[0] === "tasks" && seg.length === 2 && m === "DELETE") {
const id = seg[1];
const index = mockTasks.findIndex((item) => item.id === id);
@@ -88,6 +291,14 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
mockFindings.splice(i, 1);
findingsDeleted++;
}
} else {
// PostgreSQL uses ON DELETE SET NULL for retained findings. Keep the mock
// grouped view consistent by moving them into the unassigned/deleted bucket.
for (const finding of mockFindings) {
if (finding.task_id !== id) continue;
finding.task_id = undefined;
finding.task_description = "";
}
}
let llmRecordsDeleted = 0;
@@ -133,7 +344,56 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
reopened_intents: 0,
};
}
if (seg[0] === "tasks" && seg[2] === "control") return { id: seg[1], paused: b.action === "pause" };
if (seg[0] === "tasks" && seg[2] === "llm" && seg[3] === "resolution" && m === "GET") {
const task = mockTasks.find((item) => item.id === seg[1]);
if (!task) throw new Error("任务不存在");
return {
mainagent: mockRoleResolution(task, "mainagent"),
planner: mockRoleResolution(task, "planner"),
worker: mockRoleResolution(task, "worker"),
};
}
if (path === "/tasks/control/batch" && m === "POST") {
const action = b.action === "resume" ? "resume" : "pause";
return { items: bodyIDs(b.task_ids).map((id) => controlMockTask(id, action)) };
}
if (seg[0] === "tasks" && seg[2] === "intents" && seg[3] === "control" && seg[4] === "batch" && m === "POST") {
const action = b.action === "resume" ? "resume" : "pause";
return { items: bodyIDs(b.intent_ids).map((id) => controlMockIntent(id, action)) };
}
if (seg[0] === "tasks" && seg[2] === "intents" && seg[4] === "control" && m === "POST") {
const id = seg[3];
if (b.action === "cancel") {
const index = mockIntents.findIndex((item) => item.id === id);
const intent = mockIntents[index];
if (!intent || intent.inherited || (intent.state !== "running" && intent.state !== "paused")) {
throw new Error("Worker 状态已变化");
}
mockIntents.splice(index, 1);
return {
id: Number(id.replace(/\D/g, "")) || 0,
state: "cancelled",
deleted: {
intents: 1,
facts: 1,
findings: 1,
activities: D.activity.filter((item) => item.intent_id === id).length,
},
};
}
const action = b.action === "resume" ? "resume" : "pause";
const result = controlMockIntent(id, action);
if (!result.ok) throw new Error(result.error ?? "Worker 状态已变化");
return { id: Number(id.replace(/\D/g, "")) || 0, state: result.state };
}
if (seg[0] === "tasks" && seg.length === 3 && seg[2] === "control" && m === "POST") {
const action = b.action === "resume" ? "resume" : "pause";
const result = controlMockTask(seg[1], action);
if (!result.ok) throw new Error(result.error ?? "任务状态已变化");
const task = mockTasks.find((item) => item.id === seg[1]);
return { id: seg[1], paused: Boolean(task?.paused), queued: Boolean(task?.queued), status: task?.status ?? "" };
}
if (seg[0] === "tasks" && seg[2] === "chat" && seg[3] === "status") return { running: false };
if (seg[0] === "tasks" && seg[2] === "chat" && seg[3] === "stop") return { status: "stopped" };
if (path === "/active") {
const id = String(b.id ?? mockActiveTask);
@@ -227,21 +487,64 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
}
// ── assets ──
if (path === "/assets/counts") return D.assetCounts;
if (path === "/assets/counts") return mockAssetCounts();
if (path === "/assets" && m === "GET") {
const type = q.get("type") ?? "";
const list = type ? D.assets.filter((a) => a.type === type) : D.assets;
const list = type ? mockAssets.filter((a) => a.type === type) : mockAssets;
const limit = Number(q.get("limit") ?? 50);
const offset = Number(q.get("offset") ?? 0);
return { count: list.length, total: list.length, assets: list.slice(offset, offset + limit) };
}
if (path === "/assets" && m === "DELETE") return { deleted: (b.ids as unknown[])?.length ?? 0 };
if (path === "/assets" && m === "DELETE") {
const ids = new Set(Array.isArray(b.ids) ? b.ids.map(Number) : []);
let deleted = 0;
for (let index = mockAssets.length - 1; index >= 0; index--) {
if (!ids.has(mockAssets[index].id)) continue;
mockAssets.splice(index, 1);
deleted++;
}
return { deleted };
}
// ── companies ──
if (path === "/companies" && m === "GET") return D.companies;
if (path === "/companies" && m === "POST") return { id: 2, created: true, scope_added: 0 };
if (seg[0] === "companies" && seg[2] === "scope") return { added: 0, skipped: 0, invalid: 0 };
if (seg[0] === "companies" && seg.length === 2 && m === "DELETE") return { deleted: 1, assets_deleted: 0 };
if (path === "/companies" && m === "GET") return structuredClone(mockCompanies);
if (path === "/companies" && m === "POST") {
const name = String(b.name ?? "").trim();
if (!name) throw new Error("企业名称不能为空");
if (mockCompanies.some((company) => company.name.toLowerCase() === name.toLowerCase())) {
throw new Error("企业已存在");
}
const id = mockCompanies.reduce((max, company) => Math.max(max, company.id), 0) + 1;
const scope = mockScopeRows(id, b.scope);
const company: Company = { id, name, asset_count: 0, scope };
mockCompanies.push(company);
return { id, created: true, scope_added: scope.length, scope_invalid: 0 };
}
if (seg[0] === "companies" && seg[2] === "scope" && m === "POST") {
const company = mockCompanies.find((item) => item.id === Number(seg[1]));
if (!company) throw new Error("企业不存在");
const rows = mockScopeRows(company.id, b.scope);
company.scope = b.reset === true ? rows : [...(company.scope ?? []), ...rows];
return { added: rows.length, skipped: 0, invalid: 0 };
}
if (seg[0] === "companies" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockCompanies.findIndex((item) => item.id === id);
if (index < 0) throw new Error("企业不存在");
mockCompanies.splice(index, 1);
let assetsDeleted = 0;
for (let assetIndex = mockAssets.length - 1; assetIndex >= 0; assetIndex--) {
const asset: Asset = mockAssets[assetIndex];
if (asset.company_id !== id) continue;
if (b.delete_assets === true) {
mockAssets.splice(assetIndex, 1);
assetsDeleted++;
} else {
delete asset.company_id;
}
}
return { deleted: 1, assets_deleted: assetsDeleted };
}
// ── exploration ──
if (path === "/exploration/frontier") return D.frontier;
@@ -267,6 +570,72 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
tasks,
};
}
if (path === "/exploration/findings/groups") {
const severityOrder = { critical: 4, high: 3, medium: 2, low: 1 } as const;
let list = mockFindings.slice();
const filterSeverity = q.get("severity");
const filterStatus = q.get("status");
const filterVulnclass = q.get("vulnclass");
const filterTask = q.get("task_id");
if (filterSeverity) list = list.filter((finding) => finding.severity === filterSeverity);
if (filterStatus) list = list.filter((finding) => finding.status === filterStatus);
if (filterVulnclass) list = list.filter((finding) => finding.vulnclass === filterVulnclass);
if (filterTask === "__unassigned__") list = list.filter((finding) => !finding.task_id);
else if (filterTask) list = list.filter((finding) => finding.task_id === filterTask);
const grouped = new Map<string, typeof list>();
for (const finding of list) {
const key = finding.task_id ?? "__unassigned__";
grouped.set(key, [...(grouped.get(key) ?? []), finding]);
}
const groups = Array.from(grouped, ([key, items]) => {
const owner = mockTasks.find((candidate) => candidate.id === key);
return {
task_id: key === "__unassigned__" ? null : key,
task_description: owner?.description ?? items[0]?.task_description ?? "",
task_status: owner?.status ?? "",
count: items.length,
critical: items.filter((finding) => finding.severity === "critical").length,
high: items.filter((finding) => finding.severity === "high").length,
medium: items.filter((finding) => finding.severity === "medium").length,
low: items.filter((finding) => finding.severity === "low").length,
last_found_at: items.reduce((latest, finding) => (finding.ts > latest ? finding.ts : latest), ""),
max_severity: Math.max(...items.map((finding) => severityOrder[finding.severity])),
};
});
groups.sort((left, right) =>
q.get("sort") === "severity"
? right.max_severity - left.max_severity || right.last_found_at.localeCompare(left.last_found_at)
: right.last_found_at.localeCompare(left.last_found_at),
);
const rawPage = Number(q.get("page") ?? 1);
const rawPageSize = Number(q.get("limit") ?? 10);
const page = Number.isFinite(rawPage) && rawPage > 0 ? Math.floor(rawPage) : 1;
const pageSize = Number.isFinite(rawPageSize) && rawPageSize > 0 ? Math.min(100, Math.floor(rawPageSize)) : 10;
return {
items: groups.slice((page - 1) * pageSize, page * pageSize),
total: groups.length,
finding_total: list.length,
page,
page_size: pageSize,
};
}
if (seg[0] === "exploration" && seg[1] === "findings" && seg[3] === "deepen" && m === "POST") {
const finding = mockFindings.find((candidate) => candidate.id === seg[2]);
const description = String(b.description ?? "").trim();
if (!description) throw new Error("description is required");
if ([...description].length > 4000) throw new Error("description must be at most 4000 characters");
if (!finding) throw new Error("finding not found");
if (!finding.task_id || !mockTasks.some((candidate) => candidate.id === finding.task_id)) {
throw new Error("finding origin task or node is no longer available");
}
return {
task_id: finding.task_id,
intent_id: `mock-deepen-${Date.now()}`,
state: "open",
queued: false,
};
}
// 单条 finding:GET 详情 / PATCH 改状态/严重度/名称/类别(demo 直接改内存对象)。
if (seg[0] === "exploration" && seg[1] === "findings" && seg.length === 3 && seg[2] !== "stats") {
const f = mockFindings.find((x) => x.id === seg[2]);
@@ -320,14 +689,17 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (fSev) list = list.filter((f) => f.severity === fSev);
if (fStatus) list = list.filter((f) => f.status === fStatus);
if (fVuln) list = list.filter((f) => f.vulnclass === fVuln);
if (fTask) list = list.filter((f) => f.task_id === fTask);
if (fTask === "__unassigned__") list = list.filter((f) => !f.task_id);
else if (fTask) list = list.filter((f) => f.task_id === fTask);
list.sort((a, b) =>
q.get("sort") === "severity"
? sev[b.severity] - sev[a.severity] || +new Date(b.ts) - +new Date(a.ts)
: +new Date(b.ts) - +new Date(a.ts),
);
const page = Number(q.get("page") ?? 1);
const pageSize = Number(q.get("limit") ?? 20);
const rawPage = Number(q.get("page") ?? 1);
const rawPageSize = Number(q.get("limit") ?? 20);
const page = Number.isFinite(rawPage) && rawPage > 0 ? Math.floor(rawPage) : 1;
const pageSize = Number.isFinite(rawPageSize) && rawPageSize > 0 ? Math.min(200, Math.floor(rawPageSize)) : 20;
return {
items: list.slice((page - 1) * pageSize, page * pageSize).map(withFid),
total: list.length,
@@ -339,13 +711,20 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (q.has("page")) {
const before = Number(q.get("before") ?? 0);
const limit = Math.max(1, Number(q.get("limit") ?? 300));
let list = D.intents;
let list = mockIntents;
if (before > 0) list = list.filter((intent) => Number(intent.id.replace(/\D/g, "") || intent.id) < before);
return { items: list.slice(0, limit), has_more: list.length > limit };
}
return D.intents;
return mockIntents;
}
if (path === "/exploration/tokens") {
const selectedTask = task ? mockTasks.find((item) => item.id === task) : undefined;
return {
workers: D.tokenWorkers,
sessions: D.tokenSessions,
total: selectedTask?.tokens ?? D.tokenTotal,
};
}
if (path === "/exploration/tokens") return { workers: D.tokenWorkers, total: D.tokenTotal };
if (path === "/exploration/graph") return D.explorationGraph;
if (path === "/exploration/activity" && seg.length === 2) {
const since = Number(q.get("since") ?? 0);
@@ -425,15 +804,43 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (seg[0] === "agents" && seg[2] === "visibility" && m === "GET") return D.agentDetail(seg[1]).visibility;
// ── conversations ──
if (path === "/conversations" && m === "GET") return { conversations: D.conversations };
if (path === "/conversations" && m === "POST")
return {
id: 3,
if (path === "/conversations" && m === "GET") {
sortMockConversations();
return { conversations: structuredClone(mockConversations) };
}
if (path === "/conversations" && m === "POST") {
const now = new Date().toISOString();
const title = String(b.title ?? "").trim() || "新对话";
const conversation: Conversation = {
id: mockConversations.reduce((max, item) => Math.max(max, item.id), 0) + 1,
agent_key: String(b.agent_key ?? "mainagent"),
title: String(b.title ?? "新会话"),
created_at: "2026-07-26T04:00:00Z",
updated_at: "2026-07-26T04:00:00Z",
title,
llm_profile_id: typeof b.llm_profile_id === "number" ? b.llm_profile_id : undefined,
pinned: false,
created_at: now,
updated_at: now,
};
mockConversations.unshift(conversation);
return structuredClone(conversation);
}
if (seg[0] === "conversations" && seg.length === 2 && m === "PATCH") {
const conversation = mockConversations.find((item) => item.id === Number(seg[1]));
if (!conversation) return {};
if (typeof b.title === "string") conversation.title = b.title.trim();
if (typeof b.pinned === "boolean") {
conversation.pinned = b.pinned;
conversation.pinned_at = b.pinned ? (conversation.pinned_at ?? new Date().toISOString()) : null;
}
conversation.updated_at = new Date().toISOString();
sortMockConversations();
return structuredClone(conversation);
}
if (seg[0] === "conversations" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockConversations.findIndex((item) => item.id === id);
if (index >= 0) mockConversations.splice(index, 1);
return { deleted: id };
}
if (seg[0] === "conversations" && seg[2] === "messages" && seg.length === 3 && m === "GET") {
const items = D.conversationMessages[Number(seg[1])] ?? [];
return { items, cursor: items.length ? items[items.length - 1].seq : 0, running: false };
web/src/lib/types.ts
+85 -4
View File
@@ -30,6 +30,16 @@ export interface Task {
llm_failover_state?: "default" | "ready" | "chain_exhausted" | string;
llm_failover_reason?: string;
source_task_ids?: string[]; // directly related tasks inherited as read-only context
company_ids?: number[]; // associated company asset scopes available as task context
}
export interface TaskTemplate {
id: number;
name: string;
description: string;
goal: string;
created_at: string;
updated_at: string;
}
export interface DeleteTaskOptions {
@@ -237,13 +247,22 @@ export interface TaskScopeRow {
reason?: string;
}
export type CompanyScopeKind = "domain" | "ip" | "cidr" | "icp" | "keyword";
// 新增企业时提交的结构化资产范围规则。
export interface CompanyScopeRule {
kind: CompanyScopeKind;
value: string;
}
// 公司资产范围规则的一条（归属唯一真值来源）。
export interface ScopeRow {
id: number;
company_id: number;
kind: "domain" | "ip" | "cidr";
kind: CompanyScopeKind;
domain?: string; // kind=domain 时有值
net?: string; // kind=ip|cidr 时有值
value?: string; // kind=icp|keyword 时可能由后端直接返回
raw: string; // 原始用户输入，用于显示和回填
reason?: string;
}
@@ -254,13 +273,13 @@ export interface Company {
name: string;
logo?: string; // 远程图标 URL；空则前端用名称首字母
asset_count: number;
scope: ScopeRow[];
scope?: ScopeRow[];
}
// ---- Exploration graph (per task) ----
export type ExploreKind = "task" | "begin" | "goal" | "intent" | "fact" | "finding" | "hint";
export type GoalState = "open" | "met" | "abandoned";
export type IntentState = "open" | "running" | "done" | "blocked" | "exhausted" | "stopped";
export type IntentState = "open" | "running" | "paused" | "done" | "blocked" | "exhausted" | "stopped";
export type FindingState = "confirmed" | "dismissed";
export type HintState = "active" | "consumed";
export type ExploreRel = "spawns" | "derived_from" | "yields" | "proves";
@@ -326,6 +345,33 @@ export interface FindingsPage {
page_size: number;
}
export interface FindingGroup {
task_id: string | number | null;
task_description: string;
task_status: string;
count: number;
critical: number;
high: number;
medium: number;
low: number;
last_found_at: string;
}
export interface FindingGroupsPage {
items: FindingGroup[];
total: number;
finding_total: number;
page: number;
page_size: number;
}
export interface FindingDeepenResponse {
task_id: string;
intent_id: string;
state: IntentState;
queued: boolean;
}
// FindingStats 是发现全表聚合(统计卡 + 漏洞类型下拉),服务端计算,不受分页影响。
export interface FindingStats {
total: number;
@@ -341,7 +387,7 @@ export interface FindingStats {
// FindingTaskOption 是发现页「按任务」筛选下拉的一项:有漏洞的任务(描述为空表示任务已删除,
// 前端回退展示 id)及其漏洞条数。
export interface FindingTaskOption {
id: number;
id: string | number;
description: string;
count: number;
}
@@ -417,6 +463,22 @@ export interface LLMTransition {
next?: LLMAuditProfile;
}
export interface TaskLLMResolution {
profile_id?: number;
name: string;
format: string;
model: string;
source: "task_chain" | "agent_binding" | "global_profile" | "environment" | "global";
available: boolean;
reason?: string;
}
export interface TaskLLMResolutions {
mainagent: TaskLLMResolution;
planner: TaskLLMResolution;
worker: TaskLLMResolution;
}
// ---- Agent triggers (P3 调度，仅自定义 agent) ----
export interface AgentTrigger {
id: number;
@@ -444,6 +506,8 @@ export interface Conversation {
agent_key: string;
title: string;
llm_profile_id?: number;
pinned?: boolean;
pinned_at?: string | null;
created_at: string;
updated_at: string;
}
@@ -479,6 +543,23 @@ export interface TokenUsage {
cache_write_tokens: number;
}
export interface SessionTokenUsage {
session: string;
input_tokens: number;
output_tokens: number;
cache_read_tokens: number;
cache_write_tokens: number;
}
export interface BatchControlItem {
id: string;
ok: boolean;
status?: string;
state?: string;
queued?: boolean;
error?: string;
}
// Whole-task (all agents) token aggregate.
export interface TokenTotal {
input_tokens: number;