From d3183dbbd280cb6de4c48aaa8d38a37d015c238b Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 21 Aug 2026 21:01:39 +0800 Subject: [PATCH] feat(tasks): add templates findings workflows and batch controls --- CHANGELOG.md | 25 + agent/blackboard_inheritance_test.go | 78 +- agent/capture.go | 18 +- agent/capture_usage_test.go | 87 + agent/tools.go | 36 + db/asset_dsl.go | 17 +- db/assets.go | 240 ++- db/assets_test.go | 50 + db/companies.go | 566 +++++-- db/companies_test.go | 150 ++ db/company_lock_test.go | 15 + db/company_scope.go | 63 +- db/company_scope_consistency_test.go | 248 +++ db/company_scope_test.go | 197 ++- db/conversation.go | 58 +- db/conversation_pin_test.go | 69 + db/exploration.go | 221 ++- db/exploration_tokens_test.go | 130 ++ db/findings.go | 204 ++- db/findings_test.go | 298 ++++ db/intent_admission.go | 32 + db/intent_control_test.go | 285 ++++ db/schema.sql | 61 +- db/task_context.go | 35 +- db/task_queue_test.go | 66 + db/task_source_limit_test.go | 27 + db/task_templates.go | 231 +++ db/task_templates_test.go | 131 ++ db/tasks.go | 131 +- db/tasks_test.go | 80 + server/assets.go | 203 ++- server/assets_scope_test.go | 265 ++++ server/conversations.go | 70 +- server/dto.go | 15 +- server/engine.go | 169 +- server/engine_cancelcause_test.go | 114 +- server/engine_timeout.go | 19 +- server/findings_groups.go | 166 ++ server/findings_groups_test.go | 321 ++++ server/goals.go | 288 +++- server/manager.go | 427 ++++- server/manager_lifecycle_test.go | 109 ++ server/orchestration.go | 15 +- server/server.go | 252 +-- server/server_mgmt.go | 44 +- server/task_admission_test.go | 790 +++++++++ server/task_control.go | 314 ++++ server/task_delete_barrier_test.go | 148 +- server/task_llm.go | 2 +- server/task_llm_test.go | 22 +- server/task_resolution.go | 123 ++ server/task_templates.go | 172 ++ server/task_templates_test.go | 163 ++ web/src/app/(main)/chat/page.tsx | 265 +++- web/src/app/(main)/function/assets/page.tsx | 920 ++++++++--- web/src/app/(main)/function/findings/page.tsx | 1406 ++++++++++------- .../tasks/detail/_tabs/sessions-tab.tsx | 1111 ++++++++----- web/src/app/(main)/function/tasks/page.tsx | 858 ++++++---- web/src/components/table-pagination.tsx | 51 +- web/src/components/task-template-controls.tsx | 455 ++++++ web/src/components/ui/alert.tsx | 6 +- web/src/lib/api.ts | 77 +- web/src/lib/mock/data.ts | 52 +- web/src/lib/mock/handler.ts | 451 +++++- web/src/lib/types.ts | 89 +- 65 files changed, 11603 insertions(+), 2198 deletions(-) create mode 100644 agent/capture_usage_test.go create mode 100644 db/company_lock_test.go create mode 100644 db/company_scope_consistency_test.go create mode 100644 db/conversation_pin_test.go create mode 100644 db/exploration_tokens_test.go create mode 100644 db/intent_admission.go create mode 100644 db/intent_control_test.go create mode 100644 db/task_queue_test.go create mode 100644 db/task_templates.go create mode 100644 db/task_templates_test.go create mode 100644 server/assets_scope_test.go create mode 100644 server/findings_groups.go create mode 100644 server/findings_groups_test.go create mode 100644 server/manager_lifecycle_test.go create mode 100644 server/task_admission_test.go create mode 100644 server/task_control.go create mode 100644 server/task_resolution.go create mode 100644 server/task_templates.go create mode 100644 server/task_templates_test.go create mode 100644 web/src/components/task-template-controls.tsx diff --git a/CHANGELOG.md b/CHANGELOG.md index 40928b0..83cd8ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/agent/blackboard_inheritance_test.go b/agent/blackboard_inheritance_test.go index 978c8bf..343d422 100644 --- a/agent/blackboard_inheritance_test.go +++ b/agent/blackboard_inheritance_test.go @@ -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 { diff --git a/agent/capture.go b/agent/capture.go index ca36acc..965a9f3 100644 --- a/agent/capture.go +++ b/agent/capture.go @@ -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 diff --git a/agent/capture_usage_test.go b/agent/capture_usage_test.go new file mode 100644 index 0000000..426451d --- /dev/null +++ b/agent/capture_usage_test.go @@ -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) + } +} diff --git a/agent/tools.go b/agent/tools.go index d0958a2..05b1dcc 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -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) } diff --git a/db/asset_dsl.go b/db/asset_dsl.go index 24c7f25..dbe8683 100644 --- a/db/asset_dsl.go +++ b/db/asset_dsl.go @@ -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 diff --git a/db/assets.go b/db/assets.go index c1bcb47..61be527 100644 --- a/db/assets.go +++ b/db/assets.go @@ -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)) diff --git a/db/assets_test.go b/db/assets_test.go index 6e92030..c2ae2e5 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -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) + } +} diff --git a/db/companies.go b/db/companies.go index 84be395..8e73acc 100644 --- a/db/companies.go +++ b/db/companies.go @@ -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 } diff --git a/db/companies_test.go b/db/companies_test.go index 091e153..395c012 100644 --- a/db/companies_test.go +++ b/db/companies_test.go @@ -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 { diff --git a/db/company_lock_test.go b/db/company_lock_test.go new file mode 100644 index 0000000..5f4ea24 --- /dev/null +++ b/db/company_lock_test.go @@ -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) + } + } +} diff --git a/db/company_scope.go b/db/company_scope.go index 8ed28f0..df2b5fb 100644 --- a/db/company_scope.go +++ b/db/company_scope.go @@ -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. diff --git a/db/company_scope_consistency_test.go b/db/company_scope_consistency_test.go new file mode 100644 index 0000000..19e83fd --- /dev/null +++ b/db/company_scope_consistency_test.go @@ -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) + } +} diff --git a/db/company_scope_test.go b/db/company_scope_test.go index ec5f86d..2139e7d 100644 --- a/db/company_scope_test.go +++ b/db/company_scope_test.go @@ -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 diff --git a/db/conversation.go b/db/conversation.go index 2b2ad41..5a7edc5 100644 --- a/db/conversation.go +++ b/db/conversation.go @@ -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 } diff --git a/db/conversation_pin_test.go b/db/conversation_pin_test.go new file mode 100644 index 0000000..da7a106 --- /dev/null +++ b/db/conversation_pin_test.go @@ -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) + } +} diff --git a/db/exploration.go b/db/exploration.go index dbcdb35..e86bf19 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -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) { diff --git a/db/exploration_tokens_test.go b/db/exploration_tokens_test.go new file mode 100644 index 0000000..a0bbf16 --- /dev/null +++ b/db/exploration_tokens_test.go @@ -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) + } +} diff --git a/db/findings.go b/db/findings.go index 48f505c..485a016 100644 --- a/db/findings.go +++ b/db/findings.go @@ -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 导出(导出当前筛选/全部)。结果按严重等级降序、 diff --git a/db/findings_test.go b/db/findings_test.go index 4d67c09..4ed3ce2 100644 --- a/db/findings_test.go +++ b/db/findings_test.go @@ -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) + } +} diff --git a/db/intent_admission.go b/db/intent_admission.go new file mode 100644 index 0000000..5c7b0dc --- /dev/null +++ b/db/intent_admission.go @@ -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() +} diff --git a/db/intent_control_test.go b/db/intent_control_test.go new file mode 100644 index 0000000..6cfb9cb --- /dev/null +++ b/db/intent_control_test.go @@ -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) + } + } +} diff --git a/db/schema.sql b/db/schema.sql index fd55259..a076156 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -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(); diff --git a/db/task_context.go b/db/task_context.go index 4927b02..14053d4 100644 --- a/db/task_context.go +++ b/db/task_context.go @@ -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: diff --git a/db/task_queue_test.go b/db/task_queue_test.go new file mode 100644 index 0000000..a9b949b --- /dev/null +++ b/db/task_queue_test.go @@ -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) + } +} diff --git a/db/task_source_limit_test.go b/db/task_source_limit_test.go index 01596c8..7c54022 100644 --- a/db/task_source_limit_test.go +++ b/db/task_source_limit_test.go @@ -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) + } +} diff --git a/db/task_templates.go b/db/task_templates.go new file mode 100644 index 0000000..2a5865a --- /dev/null +++ b/db/task_templates.go @@ -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 +} diff --git a/db/task_templates_test.go b/db/task_templates_test.go new file mode 100644 index 0000000..a3425b9 --- /dev/null +++ b/db/task_templates_test.go @@ -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) + } +} diff --git a/db/tasks.go b/db/tasks.go index 3740ccc..3af6275 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -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 diff --git a/db/tasks_test.go b/db/tasks_test.go index 655b06b..0f622e3 100644 --- a/db/tasks_test.go +++ b/db/tasks_test.go @@ -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) + } +} diff --git a/server/assets.go b/server/assets.go index ac8dcba..5c31820 100644 --- a/server/assets.go +++ b/server/assets.go @@ -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{ diff --git a/server/assets_scope_test.go b/server/assets_scope_test.go new file mode 100644 index 0000000..c3e6b7c --- /dev/null +++ b/server/assets_scope_test.go @@ -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()) + } +} diff --git a/server/conversations.go b/server/conversations.go index 69294b8..de81e4c 100644 --- a/server/conversations.go +++ b/server/conversations.go @@ -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) { diff --git a/server/dto.go b/server/dto.go index 599e47f..87b8ba1 100644 --- a/server/dto.go +++ b/server/dto.go @@ -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, } } diff --git a/server/engine.go b/server/engine.go index 4606a8f..551aa0a 100644 --- a/server/engine.go +++ b/server/engine.go @@ -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(): diff --git a/server/engine_cancelcause_test.go b/server/engine_cancelcause_test.go index 5501821..d1f4f7a 100644 --- a/server/engine_cancelcause_test.go +++ b/server/engine_cancelcause_test.go @@ -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") + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 7d98e27..8d57293 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -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) { diff --git a/server/findings_groups.go b/server/findings_groups.go new file mode 100644 index 0000000..86f0972 --- /dev/null +++ b/server/findings_groups.go @@ -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, + }) +} diff --git a/server/findings_groups_test.go b/server/findings_groups_test.go new file mode 100644 index 0000000..5cb44ed --- /dev/null +++ b/server/findings_groups_test.go @@ -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) + } +} diff --git a/server/goals.go b/server/goals.go index 3d83708..dbc497d 100644 --- a/server/goals.go +++ b/server/goals.go @@ -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) } } diff --git a/server/manager.go b/server/manager.go index 7779088..f606652 100644 --- a/server/manager.go +++ b/server/manager.go @@ -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 = "" diff --git a/server/manager_lifecycle_test.go b/server/manager_lifecycle_test.go new file mode 100644 index 0000000..11be769 --- /dev/null +++ b/server/manager_lifecycle_test.go @@ -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) + } +} diff --git a/server/orchestration.go b/server/orchestration.go index 9ed4446..0f52a9c 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -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" diff --git a/server/server.go b/server/server.go index 0b82c82..15fd08f 100644 --- a/server/server.go +++ b/server/server.go @@ -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, ©) + } + } 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. diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 8782d56..49c0998 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -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) } diff --git a/server/task_admission_test.go b/server/task_admission_test.go new file mode 100644 index 0000000..eff7446 --- /dev/null +++ b/server/task_admission_test.go @@ -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) + } +} diff --git a/server/task_control.go b/server/task_control.go new file mode 100644 index 0000000..3811233 --- /dev/null +++ b/server/task_control.go @@ -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}) +} diff --git a/server/task_delete_barrier_test.go b/server/task_delete_barrier_test.go index 4b95192..361b161 100644 --- a/server/task_delete_barrier_test.go +++ b/server/task_delete_barrier_test.go @@ -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"}}} diff --git a/server/task_llm.go b/server/task_llm.go index c1dbe88..aa7be40 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -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() } diff --git a/server/task_llm_test.go b/server/task_llm_test.go index faf1cfd..a92d760 100644 --- a/server/task_llm_test.go +++ b/server/task_llm_test.go @@ -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(¤t, ¤t, []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) diff --git a/server/task_resolution.go b/server/task_resolution.go new file mode 100644 index 0000000..9c7d3d2 --- /dev/null +++ b/server/task_resolution.go @@ -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) +} diff --git a/server/task_templates.go b/server/task_templates.go new file mode 100644 index 0000000..2e728b1 --- /dev/null +++ b/server/task_templates.go @@ -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 +} diff --git a/server/task_templates_test.go b/server/task_templates_test.go new file mode 100644 index 0000000..a3a9b1b --- /dev/null +++ b/server/task_templates_test.go @@ -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()) + } +} diff --git a/web/src/app/(main)/chat/page.tsx b/web/src/app/(main)/chat/page.tsx index 32ecdb2..d059be2 100644 --- a/web/src/app/(main)/chat/page.tsx +++ b/web/src/app/(main)/chat/page.tsx @@ -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) { // 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({ ))} )} -
- {leftSlot} +
+ {leftSlot ?
{leftSlot}
: null} {onPickFiles && ( <> )}