diff --git a/CHANGELOG.md b/CHANGELOG.md index 73413e0..f604033 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,12 +33,14 @@ - 运行中和暂停任务支持编辑完整配置链、调整顺序和手动切换当前配置，从下一次 LLM 调用生效。 - 自动切换、手动切换和全链耗尽写入结构化系统活动，并通过任务活动流弹出可去重提醒。 - 增加任务角色模型解析接口，统一解析 Main Agent、Planner 和 Worker 下一次调用使用的配置与模型。 +- 增加可选的全局 LLM Pool，支持配置调用顺序、指定模型失败兜底、健康状态、冷却恢复和手动重置。 +- 增加始终开启的 LLM 用量账本，按任务、会话、模型和配置聚合输入、输出及缓存 Token。 #### 修改的功能 - 目标拆解、Planner、Worker 和主 Agent 共享任务级 LLM runtime（解析优先级见下）。 - 故障转移只响应明确的额度、余额或账单错误；普通限流、鉴权、网络、服务端和上下文错误不会错误切换。 -- LLM 配置页改为配置卡片列表和右侧抽屉编辑，保留任务级配置链，不引入 Pool 轮询、优先级和熔断界面。 +- LLM 配置页改为配置卡片列表和右侧抽屉编辑，并提供 Pool 轮询顺序、优先级、排除项、健康状态和恢复操作。 - 配置链耗尽信息支持窄屏换行；当前模型由会话列表图标改为当前会话标题旁的文字标签和完整 Tooltip。 - 任务内各角色解析 LLM 的顺序调整为「Agent 绑定 → 任务配置链 → 全局」：显式绑定模型的角色始终跑在该模型上，未绑定角色才落到任务链，任务链为空时再落到全局。 - 出站代理留空即直连，不再回退 `HTTP_PROXY`/`HTTPS_PROXY` 环境变量（显式 `ARTEX_LLM_PROXY` 不受影响）；代理输入支持带账号密码的 `socks5://user:pass@host:port`。 @@ -55,13 +57,16 @@ #### 新增的功能 +- 任务支持全局分类的创建、重命名、删除和筛选；新建任务可直接选择分类。 - 增加全局任务模板 CRUD 和右侧管理抽屉；新建任务可加载预设描述与目标，也可将当前内容另存为模板。 - 新建任务支持关联多个来源任务和多个企业资产范围；企业范围通过单个多行文本框自动识别域名、URL、IP、CIDR、ICP 和企业关键词。 +- 任务测试资产支持实时新增和删除，记录人工、企业、继承或 Agent 发现等来源；Worker 会话标题旁显示当前测试资产及来源摘要。 - 发现页按任务分组并提供组内独立分页，支持为漏洞填写描述并创建高优先级深入意图。 - 对话支持重命名、置顶、取消置顶和删除；任务列表支持当前页批量暂停与继续。 - 流量详情和工具执行详情改为右侧抽屉；HTTP 报文补齐 `Host` 并高亮请求行、状态码、Header、JSON 和标记语言正文。 - 任务列表操作列新增「详情」与「暂停/继续」按钮；会话发送键位可在系统设置中配置（`Enter` / `Cmd+Enter` 等），Web 搜索区新增代理输入框。 - 会话标题右侧徽标改为显示 LLM 配置名称，模型 ID 移至 hover Tooltip；任务可设置可选名称，任务列表新增名称列与搜索（空则回退到描述）。 +- 技能页增加使用次数、最近使用时间和缺失依赖统计，便于定位未生效或未安装的技能。 #### 修改的功能 @@ -136,6 +141,7 @@ - 创建任务新增「资产覆盖度功能」开关（默认开）：关闭后不计算与展示覆盖度、态势图只显示资产、不自动累积 `task_scope`，相关工具从 Planner / 主 Agent 剔除；企业关联不受此开关影响。 - `insert_assets` 每条资产新增 `related` 标记（默认 `true`）：仅在开启覆盖度时生效，`false` 只入共享资产库、不计入本任务覆盖度（如顺带发现的旁站或无关资产）。 +- 企业范围和任务测试资产统一支持域名、URL、IP、CIDR、ICP 与关键词文本识别；域名/IP 可创建或复用全局资产，CIDR/ICP/关键词保留为任务范围上下文。 ### 构建 diff --git a/agent/tools.go b/agent/tools.go index 4002889..a1830f2 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -486,6 +486,8 @@ func (t *ToolSet) graphOverviewData() map[string]any { e["value"] = r.Domain case r.Net != "": e["value"] = r.Net + case r.Value != "": + e["value"] = r.Value case r.CompanyID != nil: e["company_id"] = *r.CompanyID if t.cs != nil { @@ -804,6 +806,8 @@ func (t *ToolSet) relatedTaskOverviews() []map[string]any { entry["value"] = budget.take(row.Domain, 400) case row.Net != "": entry["value"] = budget.take(row.Net, 400) + case row.Value != "": + entry["value"] = budget.take(row.Value, 400) case row.CompanyID != nil: entry["company_id"] = *row.CompanyID } diff --git a/agent/tools_insert.go b/agent/tools_insert.go index e0b194a..5b4d58c 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "fmt" "strconv" "strings" @@ -284,6 +285,18 @@ func (t *ToolSet) insertAssets() actool.CoreTool { results = append(results, result{Index: i, ID: id, Type: typ}) t.writes.Assets++ t.anchorOwner(id) + if taskID > 0 { + var sourceNodeID *int64 + if t.ownerNode > 0 { + nodeID := t.ownerNode + sourceNodeID = &nodeID + } + summary := "Agent 通过 insert_assets 登记" + if t.ownerNode > 0 { + summary = fmt.Sprintf("Worker 意图 #%d 通过 insert_assets 登记", t.ownerNode) + } + _ = t.as.SetTaskAssetSource(taskID, id, "agent", summary, sourceNodeID) + } // 自动入测试范围(source='auto')：只对 worker 顶层显式插入的这一项，按其 // 类型加保守范围；side-effect 派生的资产不经此处，故范围不盲目扩大。taskID=0 时无操作。 // 资产覆盖度功能关闭时不再累积测试范围(分母)；related=false(与当前任务无关) @@ -365,14 +378,14 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { return writeTool( "add_task_scope", "把测试范围加入【本任务】——这是资产测试覆盖度的分母，也是本任务的授权边界。\n"+ - "kind 支持：company(整个公司名下资产) / root_domain(整个根域，含所有子域) / subdomain(单个精确子域) / ip / cidr。\n"+ + "kind 支持：company(整个公司名下资产) / root_domain(整个根域，含所有子域) / subdomain(单个精确子域) / ip / cidr / icp / keyword。\n"+ "说明：worker 逐个碰到的主机会被系统【自动】加进范围(精确子域)；本工具用于【主动扩大】——把整个根域/整个公司纳入，或补充指定某子域/IP。\n"+ - "value：company 传公司名或 id(公司须已存在)；root_domain/subdomain 传域名；ip/cidr 传 IP 或网段。\n"+ + "value：company 传公司名或 id(公司须已存在)；root_domain/subdomain 传域名；ip/cidr 传 IP 或网段；icp/keyword 传备案号或企业关键词。\n"+ "务必给 reason 说明依据(可审计)。多条用 entries 数组。", obj(map[string]any{ - "entries": map[string]any{"type": "array", "description": "批量：[{kind, value}]。kind∈company/root_domain/subdomain/ip/cidr；value=公司名或id / 域名 / IP / CIDR。", "items": map[string]any{"type": "object"}}, - "kind": str("[单条] company / root_domain / subdomain / ip / cidr"), - "value": str("[单条] 公司名或id / 域名 / IP / CIDR"), + "entries": map[string]any{"type": "array", "description": "批量：[{kind, value}]。kind∈company/root_domain/subdomain/ip/cidr/icp/keyword。", "items": map[string]any{"type": "object"}}, + "kind": str("[单条] company / root_domain / subdomain / ip / cidr / icp / keyword"), + "value": str("[单条] 公司名或id / 域名 / IP / CIDR / ICP / 关键词"), "reason": str("加入依据(用于审计)，务必填写"), }), func(_ context.Context, in json.RawMessage) (actool.Result, error) { @@ -404,7 +417,7 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { errs[strconv.Itoa(i)] = err.Error() continue } - added = append(added, map[string]any{"kind": ts.Kind, "domain": ts.Domain, "net": ts.Net, "company_id": ts.CompanyID}) + added = append(added, map[string]any{"kind": ts.Kind, "domain": ts.Domain, "net": ts.Net, "value": ts.Value, "company_id": ts.CompanyID}) } out := map[string]any{"added": added} if len(errs) > 0 { diff --git a/db/assets.go b/db/assets.go index 61be527..ec3a7cb 100644 --- a/db/assets.go +++ b/db/assets.go @@ -52,8 +52,11 @@ type Asset struct { Method string `json:"method,omitempty"` Params []map[string]any `json:"params,omitempty"` // meta - Extra map[string]any `json:"extra,omitempty"` - LastSeen string `json:"last_seen"` + Extra map[string]any `json:"extra,omitempty"` + LastSeen string `json:"last_seen"` + TaskSource string `json:"task_source,omitempty"` + TaskSourceSummary string `json:"task_source_summary,omitempty"` + TaskSourceNodeID *int64 `json:"task_source_node_id,omitempty"` } // AuthItem is one entry in the auth array. @@ -1140,7 +1143,14 @@ FROM assets WHERE $1 = ANY(task_ids)` return nil, err } defer rows.Close() - return scanAssets(rows) + assets, err := scanAssets(rows) + if err != nil { + return nil, err + } + if err := s.hydrateTaskAssetSources(taskID, assets); err != nil { + return nil, err + } + return assets, nil } func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) { diff --git a/db/schema.sql b/db/schema.sql index 5c517b6..59b25ec 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -337,9 +337,24 @@ CREATE TABLE IF NOT EXISTS llm_profile_health ( -- ===================================================================== -- D. 任务层 -- ===================================================================== +-- Global task categories are intentionally independent from task templates. +-- Deleting a category only moves its tasks back to the uncategorized bucket. +CREATE TABLE IF NOT EXISTS task_categories ( + id BIGSERIAL PRIMARY KEY, + name TEXT NOT NULL, + nkey TEXT NOT NULL UNIQUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_task_categories_name ON task_categories(name, id); +DROP TRIGGER IF EXISTS trg_task_categories_upd ON task_categories; +CREATE TRIGGER trg_task_categories_upd BEFORE UPDATE ON task_categories + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + CREATE TABLE IF NOT EXISTS tasks ( id BIGSERIAL PRIMARY KEY, name TEXT NOT NULL DEFAULT '', + category_id BIGINT REFERENCES task_categories(id) ON DELETE SET NULL, description TEXT NOT NULL, goal TEXT NOT NULL, exploration_id BIGINT NOT NULL UNIQUE @@ -385,8 +400,11 @@ ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued_at TIMESTAMPTZ; ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queue_mode TEXT NOT NULL DEFAULT ''; -- 可选的任务名称;补旧库。空串=未命名,前端展示时回退到描述。 ALTER TABLE tasks ADD COLUMN IF NOT EXISTS name TEXT NOT NULL DEFAULT ''; +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS category_id BIGINT REFERENCES task_categories(id) ON DELETE SET NULL; 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; +CREATE INDEX IF NOT EXISTS idx_tasks_category ON tasks(category_id, created_at DESC) + WHERE deleted_at IS NULL AND category_id IS NOT NULL; -- Reusable task description/goal presets. nkey is the normalized, case-insensitive -- identity used to reject visually equivalent duplicate names. @@ -415,6 +433,55 @@ CREATE TABLE IF NOT EXISTS task_relations ( ); CREATE INDEX IF NOT EXISTS idx_task_relations_source ON task_relations(source_task_id); +-- Task/asset provenance supplements the legacy assets.task_ids association. The +-- array remains the compatibility source for existing query and cleanup paths; +-- this relation records how each association was obtained for operator review. +CREATE TABLE IF NOT EXISTS task_asset_links ( + task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + asset_id BIGINT NOT NULL REFERENCES assets(id) ON DELETE CASCADE, + source TEXT NOT NULL DEFAULT 'system', + source_summary TEXT NOT NULL DEFAULT '', + source_node_id BIGINT REFERENCES exploration_nodes(id) ON DELETE SET NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (task_id, asset_id) +); +CREATE INDEX IF NOT EXISTS idx_task_asset_links_asset ON task_asset_links(asset_id, task_id); +CREATE INDEX IF NOT EXISTS idx_task_asset_links_node ON task_asset_links(source_node_id) + WHERE source_node_id IS NOT NULL; +DROP TRIGGER IF EXISTS trg_task_asset_links_upd ON task_asset_links; +CREATE TRIGGER trg_task_asset_links_upd BEFORE UPDATE ON task_asset_links + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + +-- Keep provenance rows synchronized when existing asset upsert paths append or +-- remove task ids. Detailed callers overwrite the generic source after upsert. +CREATE OR REPLACE FUNCTION sync_task_asset_links() RETURNS trigger AS $$ +BEGIN + INSERT INTO task_asset_links(task_id, asset_id, source, source_summary) + SELECT task.id, NEW.id, 'system', '任务执行期间自动关联' + FROM unnest(NEW.task_ids) AS requested(task_id) + JOIN tasks task ON task.id=requested.task_id AND task.deleted_at IS NULL + ON CONFLICT (task_id, asset_id) DO NOTHING; + + DELETE FROM task_asset_links link + WHERE link.asset_id=NEW.id + AND NOT (link.task_id=ANY(NEW.task_ids)); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; +DROP TRIGGER IF EXISTS trg_assets_task_links ON assets; +CREATE TRIGGER trg_assets_task_links AFTER INSERT OR UPDATE OF task_ids ON assets + FOR EACH ROW EXECUTE FUNCTION sync_task_asset_links(); + +-- Existing installations receive an auditable legacy source without rewriting +-- task_ids. Ignore stale array ids that no longer resolve to a live task. +INSERT INTO task_asset_links(task_id, asset_id, source, source_summary) +SELECT task.id, asset.id, 'legacy', '由历史任务资产关联迁移' +FROM assets asset +CROSS JOIN LATERAL unnest(asset.task_ids) AS requested(task_id) +JOIN tasks task ON task.id=requested.task_id AND task.deleted_at IS NULL +ON CONFLICT (task_id, asset_id) DO NOTHING; + -- Ordered task-level LLM failover chain. A quota-exhausted entry is skipped -- until the user saves/resets the chain, which clears all failure state. CREATE TABLE IF NOT EXISTS task_llm_profiles ( @@ -462,17 +529,24 @@ WHERE t.active_llm_profile_id IS NULL CREATE TABLE IF NOT EXISTS task_scope ( id BIGSERIAL PRIMARY KEY, task_id BIGINT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, - kind TEXT NOT NULL CHECK (kind IN ('company','root_domain','subdomain','ip','cidr')), + kind TEXT NOT NULL CHECK (kind IN ('company','root_domain','subdomain','ip','cidr','icp','keyword')), company_id BIGINT REFERENCES companies(id) ON DELETE CASCADE, -- kind='company' domain TEXT, -- root_domain / subdomain net CIDR, -- ip / cidr + value TEXT, -- icp / keyword source TEXT NOT NULL DEFAULT 'auto' CHECK (source IN ('auto','agent','manual')), reason TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +-- 旧库升级：扩展任务范围，使其与企业范围的单文本框识别能力一致。 +ALTER TABLE task_scope ADD COLUMN IF NOT EXISTS value TEXT; +ALTER TABLE task_scope DROP CONSTRAINT IF EXISTS task_scope_kind_check; +ALTER TABLE task_scope ADD CONSTRAINT task_scope_kind_check + CHECK (kind IN ('company','root_domain','subdomain','ip','cidr','icp','keyword')); -- 去重：同一 task 的同一条范围只存一次（自动填批量插入靠它幂等）。 -CREATE UNIQUE INDEX IF NOT EXISTS uq_task_scope ON task_scope( - task_id, kind, COALESCE(domain,''), COALESCE(net::text,''), COALESCE(company_id,0)); +DROP INDEX IF EXISTS uq_task_scope; +CREATE UNIQUE INDEX IF NOT EXISTS uq_task_scope_v2 ON task_scope( + task_id, kind, COALESCE(domain,''), COALESCE(net::text,''), COALESCE(company_id,0), COALESCE(value,'')); CREATE INDEX IF NOT EXISTS idx_ts_domain ON task_scope(domain) WHERE kind IN ('root_domain','subdomain'); CREATE INDEX IF NOT EXISTS idx_ts_net ON task_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr'); CREATE INDEX IF NOT EXISTS idx_ts_company ON task_scope(company_id) WHERE kind = 'company'; diff --git a/db/task_assets.go b/db/task_assets.go new file mode 100644 index 0000000..db87fef --- /dev/null +++ b/db/task_assets.go @@ -0,0 +1,421 @@ +package db + +import ( + "database/sql" + "errors" + "fmt" + "net" + "strings" + "unicode/utf8" +) + +const ( + MaxTaskAssetMutationCount = 100 + MaxTaskAssetSummaryRunes = 500 + defaultTaskAssetSource = "system" + manualTaskScopeSummary = "用户在测试资产页手工新增" +) + +var ( + ErrTaskAssetInvalid = errors.New("invalid task asset association") + ErrTaskAssetTaskNotFound = errors.New("task not found") + ErrTaskAssetAssetNotFound = errors.New("asset not found") +) + +// TaskAssetMutation summarizes one attach request. Attached counts newly added +// associations; Existing counts requested assets that were already on the task. +type TaskAssetMutation struct { + Requested int `json:"requested"` + Attached int `json:"attached"` + Existing int `json:"existing"` +} + +// TaskAssetScopeMutation summarizes one free-form scope registration. Domain +// and IP entries create or reuse global assets; every entry also becomes an +// idempotent task_scope row. +type TaskAssetScopeMutation struct { + Requested int `json:"requested"` + AssetsLinked int `json:"assets_linked"` + AssetsExisting int `json:"assets_existing"` + ScopesAdded int `json:"scopes_added"` + ScopesExisting int `json:"scopes_existing"` +} + +// IntentAsset describes an asset explicitly anchored to a worker intent. +type IntentAsset struct { + IntentID int64 `json:"intent_id"` + AssetID int64 `json:"asset_id"` + Type string `json:"type"` + Label string `json:"label"` + Source string `json:"source"` + SourceSummary string `json:"source_summary"` + SourceNodeID *int64 `json:"source_node_id,omitempty"` + SourceTaskID int64 `json:"source_task_id"` + Inherited bool `json:"inherited"` +} + +func normalizeTaskAssetIDs(ids []int64) ([]int64, error) { + if len(ids) == 0 { + return nil, fmt.Errorf("%w: asset_ids is required", ErrTaskAssetInvalid) + } + seen := make(map[int64]struct{}, len(ids)) + normalized := make([]int64, 0, len(ids)) + for _, id := range ids { + if id <= 0 { + return nil, fmt.Errorf("%w: asset id must be positive", ErrTaskAssetInvalid) + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + normalized = append(normalized, id) + if len(normalized) > MaxTaskAssetMutationCount { + return nil, fmt.Errorf("%w: at most %d assets per request", ErrTaskAssetInvalid, MaxTaskAssetMutationCount) + } + } + return normalized, nil +} + +func normalizeTaskAssetSource(source, summary string) (string, string, error) { + source = strings.TrimSpace(strings.ToLower(source)) + if source == "" { + source = defaultTaskAssetSource + } + summary = strings.TrimSpace(summary) + if utf8.RuneCountInString(summary) > MaxTaskAssetSummaryRunes { + return "", "", fmt.Errorf("%w: source summary exceeds %d characters", ErrTaskAssetInvalid, MaxTaskAssetSummaryRunes) + } + return source, summary, nil +} + +// SetTaskAssetSource improves the generic trigger-created provenance for one +// existing task association. It never creates or deletes an asset. +func (s *AssetStore) SetTaskAssetSource(taskID, assetID int64, source, summary string, sourceNodeID *int64) error { + if taskID <= 0 || assetID <= 0 { + return fmt.Errorf("%w: task and asset ids must be positive", ErrTaskAssetInvalid) + } + source, summary, err := normalizeTaskAssetSource(source, summary) + if err != nil { + return err + } + query := ` +INSERT INTO task_asset_links(task_id, asset_id, source, source_summary, source_node_id) +SELECT task.id, asset.id, $3, $4, $5 +FROM tasks task +JOIN assets asset ON asset.id=$2 AND task.id=ANY(asset.task_ids) +WHERE task.id=$1 AND task.deleted_at IS NULL +ON CONFLICT (task_id, asset_id) DO UPDATE +SET source=EXCLUDED.source, + source_summary=EXCLUDED.source_summary, + source_node_id=COALESCE(EXCLUDED.source_node_id, task_asset_links.source_node_id)` + var result sql.Result + if s.tx != nil { + result, err = s.tx.Exec(query, taskID, assetID, source, summary, sourceNodeID) + } else { + result, err = s.db.Exec(query, taskID, assetID, source, summary, sourceNodeID) + } + if err != nil { + return err + } + rows, err := result.RowsAffected() + if err != nil { + return err + } + if rows == 0 { + return fmt.Errorf("%w: task or asset association does not exist", ErrTaskAssetInvalid) + } + return nil +} + +// RegisterTaskAssetScopes accepts the same structured scope rules as enterprise +// assets. The entire request is atomic: invalid input or any storage failure +// leaves both global assets and task scope unchanged. +func (s *AssetStore) RegisterTaskAssetScopes(taskID int64, inputs []ScopeInput) (TaskAssetScopeMutation, error) { + mutation := TaskAssetScopeMutation{Requested: len(inputs)} + if taskID <= 0 { + return mutation, fmt.Errorf("%w: task id must be positive", ErrTaskAssetInvalid) + } + if len(inputs) == 0 { + return mutation, fmt.Errorf("%w: scope is required", ErrTaskAssetInvalid) + } + if err := ValidateCompanyScopeInputBounds(inputs); err != nil { + return mutation, fmt.Errorf("%w: %v", ErrTaskAssetInvalid, err) + } + parsed := make([]ParsedScope, 0, len(inputs)) + for index, input := range inputs { + rule, err := ParseScopeInput(input) + if err != nil { + return mutation, fmt.Errorf("%w: 第 %d 条范围无效: %v", ErrTaskAssetInvalid, index+1, err) + } + parsed = append(parsed, rule) + } + if err := validateParsedScopeBounds(parsed); err != nil { + return mutation, fmt.Errorf("%w: %v", ErrTaskAssetInvalid, err) + } + + tx, err := s.db.Begin() + if err != nil { + return mutation, err + } + defer tx.Rollback() //nolint:errcheck + if err := lockCompanyScopeMutation(tx); err != nil { + return mutation, err + } + var taskExists bool + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, taskID).Scan(&taskExists); err != nil { + return mutation, err + } + if !taskExists { + return mutation, ErrTaskAssetTaskNotFound + } + + scoped := &AssetStore{db: s.db, company: s.company, tx: tx} + for _, rule := range parsed { + taskScope := TaskScope{ + TaskID: taskID, + Source: "manual", + Reason: manualTaskScopeSummary, + } + var assetID int64 + switch rule.Kind { + case "domain": + taskScope.Kind = "root_domain" + taskScope.Domain = rule.Domain + var alreadyLinked bool + err := tx.QueryRow(`SELECT id, $2=ANY(task_ids) FROM assets WHERE type='root_domain' AND domain=$1`, rule.Domain, taskID). + Scan(&assetID, &alreadyLinked) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return mutation, err + } + assetID, err = scoped.UpsertRootDomain(UpsertRootDomainReq{Domain: rule.Domain, TaskID: taskID}) + if err != nil { + return mutation, err + } + if alreadyLinked { + mutation.AssetsExisting++ + } else { + mutation.AssetsLinked++ + } + case "ip": + taskScope.Kind = "ip" + taskScope.Net = rule.Net + ip, _, parseErr := net.ParseCIDR(rule.Net) + if parseErr != nil { + return mutation, fmt.Errorf("%w: 无效 IP: %s", ErrTaskAssetInvalid, rule.Raw) + } + ipValue := ip.String() + var alreadyLinked bool + err := tx.QueryRow(`SELECT id, $2=ANY(task_ids) FROM assets WHERE type='ip' AND ip=$1`, ipValue, taskID). + Scan(&assetID, &alreadyLinked) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return mutation, err + } + assetID, err = scoped.UpsertIP(UpsertIPReq{IP: ipValue, TaskID: taskID}) + if err != nil { + return mutation, err + } + if alreadyLinked { + mutation.AssetsExisting++ + } else { + mutation.AssetsLinked++ + } + case "cidr": + taskScope.Kind = "cidr" + taskScope.Net = rule.Net + case "icp", "keyword": + taskScope.Kind = rule.Kind + taskScope.Value = rule.Value + default: + return mutation, fmt.Errorf("%w: unsupported scope kind %q", ErrTaskAssetInvalid, rule.Kind) + } + + if assetID > 0 { + if err := scoped.SetTaskAssetSource(taskID, assetID, "manual", manualTaskScopeSummary, nil); err != nil { + return mutation, err + } + } + inserted, err := scoped.upsertTaskScopeResult(taskScope) + if err != nil { + return mutation, err + } + if inserted { + mutation.ScopesAdded++ + } else { + mutation.ScopesExisting++ + } + } + if err := tx.Commit(); err != nil { + return mutation, err + } + return mutation, nil +} + +// AttachAssetsToTask associates existing global assets with one live task and +// records an operator-authored source summary. Global asset rows are retained. +func (s *AssetStore) AttachAssetsToTask(taskID int64, assetIDs []int64, sourceSummary string) (TaskAssetMutation, error) { + var mutation TaskAssetMutation + assetIDs, err := normalizeTaskAssetIDs(assetIDs) + if err != nil { + return mutation, err + } + _, sourceSummary, err = normalizeTaskAssetSource("manual", sourceSummary) + if err != nil { + return mutation, err + } + if sourceSummary == "" { + return mutation, fmt.Errorf("%w: source_summary is required", ErrTaskAssetInvalid) + } + mutation.Requested = len(assetIDs) + tx, err := s.db.Begin() + if err != nil { + return mutation, err + } + defer tx.Rollback() //nolint:errcheck + + var taskExists bool + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, taskID).Scan(&taskExists); err != nil { + return mutation, err + } + if !taskExists { + return mutation, ErrTaskAssetTaskNotFound + } + var found, existing int + if err := tx.QueryRow(` +SELECT count(*), count(*) FILTER (WHERE $1=ANY(task_ids)) +FROM assets WHERE id=ANY($2::bigint[])`, taskID, assetIDs).Scan(&found, &existing); err != nil { + return mutation, err + } + if found != len(assetIDs) { + return mutation, ErrTaskAssetAssetNotFound + } + if _, err := tx.Exec(` +UPDATE assets +SET task_ids=CASE WHEN $1=ANY(task_ids) THEN task_ids ELSE array_append(task_ids,$1) END +WHERE id=ANY($2::bigint[])`, taskID, assetIDs); err != nil { + return mutation, err + } + if _, err := tx.Exec(` +INSERT INTO task_asset_links(task_id, asset_id, source, source_summary) +SELECT $1, id, 'manual', $3 FROM assets WHERE id=ANY($2::bigint[]) +ON CONFLICT (task_id, asset_id) DO UPDATE +SET source='manual', source_summary=EXCLUDED.source_summary, source_node_id=NULL`, + taskID, assetIDs, sourceSummary); err != nil { + return mutation, err + } + mutation.Existing = existing + mutation.Attached = len(assetIDs) - existing + return mutation, tx.Commit() +} + +// DetachAssetFromTask removes only the task association. The global asset and +// exploration anchors remain available for historical blackboard auditing. +func (s *AssetStore) DetachAssetFromTask(taskID, assetID int64) (bool, error) { + if taskID <= 0 || assetID <= 0 { + return false, fmt.Errorf("%w: task and asset ids must be positive", ErrTaskAssetInvalid) + } + var detachedID int64 + err := s.db.QueryRow(` +UPDATE assets SET task_ids=array_remove(task_ids,$1) +WHERE id=$2 AND $1=ANY(task_ids) +RETURNING id`, taskID, assetID).Scan(&detachedID) + if err == sql.ErrNoRows { + return false, nil + } + if err != nil { + return false, err + } + return detachedID == assetID, nil +} + +func (s *AssetStore) hydrateTaskAssetSources(taskID int64, assets []*Asset) error { + if len(assets) == 0 { + return nil + } + ids := make([]int64, 0, len(assets)) + byID := make(map[int64]*Asset, len(assets)) + for _, asset := range assets { + ids = append(ids, asset.ID) + byID[asset.ID] = asset + } + rows, err := s.db.Query(` +SELECT asset_id, source, source_summary, source_node_id +FROM task_asset_links +WHERE task_id=$1 AND asset_id=ANY($2::bigint[])`, taskID, ids) + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var assetID int64 + var source, summary string + var sourceNodeID sql.NullInt64 + if err := rows.Scan(&assetID, &source, &summary, &sourceNodeID); err != nil { + return err + } + if asset := byID[assetID]; asset != nil { + asset.TaskSource = source + asset.TaskSourceSummary = summary + if sourceNodeID.Valid { + id := sourceNodeID.Int64 + asset.TaskSourceNodeID = &id + } + } + } + return rows.Err() +} + +// IntentAssets returns all local worker targets plus immutable targets from the +// task's direct sources. Inherited non-terminal intents remain hidden, matching +// the existing source-aware session contract. +func (s *AssetStore) IntentAssets(taskID int64) ([]IntentAsset, error) { + rows, err := s.db.Query(` +WITH context AS ( + SELECT task.id AS task_id, task.exploration_id, false AS inherited + FROM tasks task + WHERE task.id=$1 AND task.deleted_at IS NULL + UNION ALL + SELECT source.id, source.exploration_id, true + FROM task_relations relation + JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL + WHERE relation.task_id=$1 +) +SELECT intent.id, asset.id, asset.type, + CASE asset.type + WHEN 'root_domain' THEN COALESCE(asset.domain,'') + WHEN 'subdomain' THEN COALESCE(asset.domain,'') + WHEN 'ip' THEN COALESCE(asset.ip,'') + WHEN 'app' THEN COALESCE(asset.app_name,'') + WHEN 'service' THEN COALESCE(NULLIF(asset.url,''), NULLIF(concat_ws(':', COALESCE(NULLIF(asset.domain,''), NULLIF(asset.ip,'')), asset.port::text),''), NULLIF(asset.service_name,''), '#' || asset.id::text) + WHEN 'endpoint' THEN COALESCE(NULLIF(asset.url,''), '#' || asset.id::text) + ELSE '#' || asset.id::text + END, + COALESCE(link.source,'anchor'), + COALESCE(NULLIF(link.source_summary,''), '意图在黑板中锚定该资产'), + link.source_node_id, context.task_id, context.inherited +FROM context +JOIN exploration_nodes intent ON intent.exploration_id=context.exploration_id AND intent.kind='intent' +JOIN exploration_anchors anchor ON anchor.node_id=intent.id +JOIN assets asset ON asset.id=anchor.asset_id +LEFT JOIN task_asset_links link ON link.task_id=context.task_id AND link.asset_id=asset.id +WHERE NOT context.inherited OR intent.state IN ('done','blocked','exhausted','stopped') +ORDER BY context.inherited, intent.id DESC, asset.id`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []IntentAsset{} + for rows.Next() { + var asset IntentAsset + var sourceNodeID sql.NullInt64 + if err := rows.Scan(&asset.IntentID, &asset.AssetID, &asset.Type, &asset.Label, + &asset.Source, &asset.SourceSummary, &sourceNodeID, &asset.SourceTaskID, &asset.Inherited); err != nil { + return nil, err + } + if sourceNodeID.Valid { + id := sourceNodeID.Int64 + asset.SourceNodeID = &id + } + out = append(out, asset) + } + return out, rows.Err() +} diff --git a/db/task_assets_context.go b/db/task_assets_context.go index 2c5b1f2..b14c79a 100644 --- a/db/task_assets_context.go +++ b/db/task_assets_context.go @@ -29,8 +29,12 @@ target AS ( JOIN task_scope ts ON ( (ts.kind='company' AND a.company_id = ts.company_id) OR (ts.kind='root_domain' AND a.root_domain = ts.domain) - OR (ts.kind='subdomain' AND a.domain = ts.domain) - OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='subdomain' AND a.domain = ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='icp' AND ( + lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value + OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value + )) ) JOIN context_tasks ctx ON ctx.task_id=ts.task_id UNION @@ -53,7 +57,7 @@ tested AS ( func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) { rows, err := s.db.Query(`WITH `+directTaskContextCTE+` SELECT ts.id, ts.task_id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''), - COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'') + COALESCE(ts.net::text,''), COALESCE(ts.value,''), ts.source, COALESCE(ts.reason,'') FROM task_scope ts JOIN context_tasks ctx ON ctx.task_id=ts.task_id LEFT JOIN companies c ON c.id=ts.company_id @@ -66,7 +70,7 @@ ORDER BY CASE WHEN ts.task_id=$1 THEN 0 ELSE 1 END, ts.id`, taskID) for rows.Next() { var scope TaskScope var companyID int64 - if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Source, &scope.Reason); err != nil { + if err := rows.Scan(&scope.ID, &scope.TaskID, &scope.Kind, &companyID, &scope.CompanyName, &scope.Domain, &scope.Net, &scope.Value, &scope.Source, &scope.Reason); err != nil { return nil, err } if companyID > 0 { @@ -175,8 +179,12 @@ context_assets AS ( JOIN task_scope ts ON ( (ts.kind='company' AND a.company_id=ts.company_id) OR (ts.kind='root_domain' AND a.root_domain=ts.domain) - OR (ts.kind='subdomain' AND a.domain=ts.domain) - OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='subdomain' AND a.domain=ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='icp' AND ( + lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value + OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value + )) ) JOIN context_tasks ctx ON ctx.task_id=ts.task_id ) diff --git a/db/task_assets_test.go b/db/task_assets_test.go new file mode 100644 index 0000000..4c19747 --- /dev/null +++ b/db/task_assets_test.go @@ -0,0 +1,188 @@ +package db + +import ( + "errors" + "fmt" + "testing" + "time" +) + +func TestRegisterTaskAssetScopesCreatesAssetsAndPersistsTextScope(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("manual scope registration", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + suffix := time.Now().UnixNano() + domain := fmt.Sprintf("manual-scope-%d.example.test", suffix) + ip := fmt.Sprintf("203.0.%d.%d", (suffix/250)%250, suffix%250+1) + inputs := []ScopeInput{ + {Kind: "domain", Value: domain}, + {Kind: "ip", Value: ip}, + {Kind: "cidr", Value: "198.51.100.0/24"}, + {Kind: "icp", Value: " 京 ICP 备 12345678 号-1 "}, + {Kind: "keyword", Value: " Acme Security "}, + } + + first, err := d.Assets().RegisterTaskAssetScopes(task.ID, inputs) + if err != nil { + t.Fatal(err) + } + if first.Requested != 5 || first.AssetsLinked != 2 || first.AssetsExisting != 0 || first.ScopesAdded != 5 || first.ScopesExisting != 0 { + t.Fatalf("unexpected first mutation: %+v", first) + } + assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0) + if err != nil || len(assets) != 2 { + t.Fatalf("task assets=%+v err=%v", assets, err) + } + assetIDs := make([]int64, 0, len(assets)) + for _, asset := range assets { + assetIDs = append(assetIDs, asset.ID) + if asset.TaskSource != "manual" || asset.TaskSourceSummary != manualTaskScopeSummary { + t.Fatalf("unexpected task asset provenance: %+v", asset) + } + } + t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs(assetIDs) }) + scopes, err := d.Assets().ListTaskScope(task.ID) + if err != nil || len(scopes) != 5 { + t.Fatalf("task scope=%+v err=%v", scopes, err) + } + values := map[string]string{} + for _, scope := range scopes { + values[scope.Kind] = scope.Value + } + if values["icp"] != NormalizeICP(inputs[3].Value) || values["keyword"] != normalizeKeyword(inputs[4].Value) { + t.Fatalf("text scope not normalized: %+v", values) + } + + second, err := d.Assets().RegisterTaskAssetScopes(task.ID, inputs) + if err != nil { + t.Fatal(err) + } + if second.AssetsExisting != 2 || second.AssetsLinked != 0 || second.ScopesExisting != 5 || second.ScopesAdded != 0 { + t.Fatalf("unexpected idempotent mutation: %+v", second) + } + + rollbackDomain := fmt.Sprintf("rollback-%d.example.test", suffix) + _, err = d.Assets().RegisterTaskAssetScopes(task.ID, []ScopeInput{ + {Kind: "domain", Value: rollbackDomain}, + {Kind: "cidr", Value: "10.0.0.0/8"}, + }) + if !errors.Is(err, ErrTaskAssetInvalid) { + t.Fatalf("invalid CIDR error=%v", err) + } + var rollbackAssets int + if err := d.QueryRow(`SELECT count(*) FROM assets WHERE type='root_domain' AND domain=$1`, rollbackDomain).Scan(&rollbackAssets); err != nil { + t.Fatal(err) + } + if rollbackAssets != 0 { + t.Fatalf("invalid request created %d assets", rollbackAssets) + } +} + +func TestTaskAssetAttachDetachPreservesGlobalAssetAndAnchors(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("asset editing", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: fmt.Sprintf("manual-%d.example.test", time.Now().UnixNano())}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) }) + + mutation, err := d.Assets().AttachAssetsToTask(task.ID, []int64{assetID, assetID}, "授权资产清单第 3 项") + if err != nil { + t.Fatal(err) + } + if mutation.Requested != 1 || mutation.Attached != 1 || mutation.Existing != 0 { + t.Fatalf("unexpected first mutation: %+v", mutation) + } + assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0) + if err != nil || len(assets) != 1 { + t.Fatalf("task assets=%+v err=%v", assets, err) + } + if assets[0].TaskSource != "manual" || assets[0].TaskSourceSummary != "授权资产清单第 3 项" { + t.Fatalf("unexpected provenance: %+v", assets[0]) + } + + intentID, err := d.Exploration(task.ExplorationID).AddIntent(map[string]any{"summary": "test asset"}, 5, []int64{assetID}, "human") + if err != nil { + t.Fatal(err) + } + detached, err := d.Assets().DetachAssetFromTask(task.ID, assetID) + if err != nil || !detached { + t.Fatalf("detach=%v err=%v", detached, err) + } + if assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0); err != nil || len(assets) != 0 { + t.Fatalf("detached task assets=%+v err=%v", assets, err) + } + var global, anchors, links int + _ = d.QueryRow(`SELECT count(*) FROM assets WHERE id=$1`, assetID).Scan(&global) + _ = d.QueryRow(`SELECT count(*) FROM exploration_anchors WHERE node_id=$1 AND asset_id=$2`, intentID, assetID).Scan(&anchors) + _ = d.QueryRow(`SELECT count(*) FROM task_asset_links WHERE task_id=$1 AND asset_id=$2`, task.ID, assetID).Scan(&links) + if global != 1 || anchors != 1 || links != 0 { + t.Fatalf("global=%d anchors=%d links=%d", global, anchors, links) + } +} + +func TestIntentAssetsIncludesDirectSourceProvenance(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + source, err := d.CreateTask("source assets", "goal", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + current, err := d.CreateTaskWithOptions("current assets", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(current.ID); _ = d.DeleteTask(source.ID) }) + + assetID, err := d.Assets().UpsertHTTPService(UpsertHTTPServiceReq{ + URL: fmt.Sprintf("https://intent-%d.example.test", time.Now().UnixNano()), TaskID: source.ID, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) }) + intentID, err := d.Exploration(source.ExplorationID).AddIntent(map[string]any{"summary": "source worker"}, 5, []int64{assetID}, "planner") + if err != nil { + t.Fatal(err) + } + if err := d.Exploration(source.ExplorationID).SetNodeState(intentID, "done"); err != nil { + t.Fatal(err) + } + nodeID := intentID + if err := d.Assets().SetTaskAssetSource(source.ID, assetID, "agent", "Worker 通过 insert_assets 登记", &nodeID); err != nil { + t.Fatal(err) + } + + assets, err := d.Assets().IntentAssets(current.ID) + if err != nil { + t.Fatal(err) + } + if len(assets) != 1 || assets[0].IntentID != intentID || assets[0].SourceTaskID != source.ID || !assets[0].Inherited { + t.Fatalf("unexpected intent assets: %+v", assets) + } + if assets[0].Source != "agent" || assets[0].SourceSummary == "" || assets[0].SourceNodeID == nil || *assets[0].SourceNodeID != intentID { + t.Fatalf("unexpected intent provenance: %+v", assets[0]) + } +} diff --git a/db/task_categories.go b/db/task_categories.go new file mode 100644 index 0000000..5a4b30a --- /dev/null +++ b/db/task_categories.go @@ -0,0 +1,201 @@ +package db + +import ( + "database/sql" + "errors" + "fmt" + "strings" + "time" + "unicode/utf8" + + "github.com/jackc/pgx/v5/pgconn" +) + +const MaxTaskCategoryNameRunes = 80 + +var ( + ErrTaskCategoryInvalid = errors.New("invalid task category") + ErrTaskCategoryNameConflict = errors.New("task category name already exists") + ErrTaskCategoryNotFound = errors.New("task category not found") + ErrTaskCategoryTaskNotFound = errors.New("task not found") +) + +// TaskCategory is a globally reusable task grouping label. +type TaskCategory struct { + ID int64 `json:"id"` + Name string `json:"name"` + NKey string `json:"-"` + TaskCount int `json:"task_count"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +const taskCategoryCols = `category.id, category.name, category.nkey, + count(task.id) FILTER (WHERE task.deleted_at IS NULL), + category.created_at, category.updated_at` + +func scanTaskCategory(row interface{ Scan(...any) error }) (TaskCategory, error) { + var category TaskCategory + err := row.Scan(&category.ID, &category.Name, &category.NKey, &category.TaskCount, &category.CreatedAt, &category.UpdatedAt) + return category, err +} + +func normalizeTaskCategoryName(name string) (string, string, error) { + name = strings.Join(strings.Fields(name), " ") + if name == "" { + return "", "", fmt.Errorf("%w: name is required", ErrTaskCategoryInvalid) + } + if utf8.RuneCountInString(name) > MaxTaskCategoryNameRunes { + return "", "", fmt.Errorf("%w: name exceeds %d characters", ErrTaskCategoryInvalid, MaxTaskCategoryNameRunes) + } + return name, strings.ToLower(name), nil +} + +func taskCategoryUniqueViolation(err error) bool { + var pgErr *pgconn.PgError + return errors.As(err, &pgErr) && pgErr.Code == "23505" +} + +func (d *DB) CreateTaskCategory(name string) (*TaskCategory, error) { + name, nkey, err := normalizeTaskCategoryName(name) + if err != nil { + return nil, err + } + category, err := scanTaskCategory(d.QueryRow(` +WITH inserted AS ( + INSERT INTO task_categories(name, nkey) + VALUES ($1,$2) + ON CONFLICT (nkey) DO NOTHING + RETURNING * +) +SELECT inserted.id, inserted.name, inserted.nkey, 0, inserted.created_at, inserted.updated_at +FROM inserted`, name, nkey)) + if err == sql.ErrNoRows { + return nil, ErrTaskCategoryNameConflict + } + if err != nil { + return nil, err + } + return &category, nil +} + +func (d *DB) ListTaskCategories() ([]*TaskCategory, error) { + rows, err := d.Query(` +SELECT ` + taskCategoryCols + ` +FROM task_categories category +LEFT JOIN tasks task ON task.category_id=category.id +GROUP BY category.id +ORDER BY category.name, category.id`) + if err != nil { + return nil, err + } + defer rows.Close() + categories := []*TaskCategory{} + for rows.Next() { + category, err := scanTaskCategory(rows) + if err != nil { + return nil, err + } + categories = append(categories, &category) + } + return categories, rows.Err() +} + +func (d *DB) GetTaskCategory(id int64) (*TaskCategory, error) { + category, err := scanTaskCategory(d.QueryRow(` +SELECT `+taskCategoryCols+` +FROM task_categories category +LEFT JOIN tasks task ON task.category_id=category.id +WHERE category.id=$1 +GROUP BY category.id`, id)) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + return &category, nil +} + +func (d *DB) RenameTaskCategory(id int64, name string) (*TaskCategory, error) { + name, nkey, err := normalizeTaskCategoryName(name) + if err != nil { + return nil, err + } + category, err := scanTaskCategory(d.QueryRow(` +WITH updated AS ( + UPDATE task_categories SET name=$2, nkey=$3 WHERE id=$1 RETURNING * +) +SELECT updated.id, updated.name, updated.nkey, + (SELECT count(*) FROM tasks WHERE category_id=updated.id AND deleted_at IS NULL), + updated.created_at, updated.updated_at +FROM updated`, id, name, nkey)) + if err == sql.ErrNoRows { + return nil, ErrTaskCategoryNotFound + } + if taskCategoryUniqueViolation(err) { + return nil, ErrTaskCategoryNameConflict + } + if err != nil { + return nil, err + } + return &category, nil +} + +// DeleteTaskCategory moves affected tasks to the uncategorized bucket through +// the tasks.category_id ON DELETE SET NULL foreign key. +func (d *DB) DeleteTaskCategory(id int64) (bool, error) { + result, err := d.Exec(`DELETE FROM task_categories WHERE id=$1`, id) + if err != nil { + return false, err + } + rows, err := result.RowsAffected() + return rows > 0, err +} + +// SetTaskCategory updates one live task. A nil category means uncategorized. +func (d *DB) SetTaskCategory(taskID int64, categoryID *int64) (*TaskCategory, error) { + if categoryID == nil { + result, err := d.Exec(`UPDATE tasks SET category_id=NULL WHERE id=$1 AND deleted_at IS NULL`, taskID) + if err != nil { + return nil, err + } + rows, err := result.RowsAffected() + if err != nil { + return nil, err + } + if rows == 0 { + return nil, ErrTaskCategoryTaskNotFound + } + return nil, nil + } + if *categoryID <= 0 { + return nil, fmt.Errorf("%w: category id must be positive", ErrTaskCategoryInvalid) + } + category, err := scanTaskCategory(d.QueryRow(` +WITH selected AS ( + SELECT * FROM task_categories WHERE id=$2 +), updated AS ( + UPDATE tasks SET category_id=$2 + WHERE id=$1 AND deleted_at IS NULL AND EXISTS (SELECT 1 FROM selected) + RETURNING id +) +SELECT selected.id, selected.name, selected.nkey, + (SELECT count(*) FROM tasks WHERE category_id=selected.id AND deleted_at IS NULL), + selected.created_at, selected.updated_at +FROM selected, updated`, taskID, *categoryID)) + if err == sql.ErrNoRows { + var taskExists bool + if checkErr := d.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, taskID).Scan(&taskExists); checkErr != nil { + return nil, checkErr + } + if !taskExists { + return nil, ErrTaskCategoryTaskNotFound + } + return nil, ErrTaskCategoryNotFound + } + if err != nil { + return nil, err + } + return &category, nil +} diff --git a/db/task_categories_test.go b/db/task_categories_test.go new file mode 100644 index 0000000..e1d5d60 --- /dev/null +++ b/db/task_categories_test.go @@ -0,0 +1,91 @@ +package db + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func TestTaskCategoryCRUDAndTaskAssignment(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + name := fmt.Sprintf("Category %d", time.Now().UnixNano()) + category, err := d.CreateTaskCategory(" " + strings.ReplaceAll(name, " ", " ") + " ") + if err != nil { + t.Fatal(err) + } + if category.Name != name { + t.Fatalf("category name=%q, want %q", category.Name, name) + } + if _, err := d.CreateTaskCategory(strings.ToUpper(name)); !errors.Is(err, ErrTaskCategoryNameConflict) { + t.Fatalf("duplicate error=%v, want %v", err, ErrTaskCategoryNameConflict) + } + + task, err := d.CreateTaskWithOptions("categorized task", "goal", TaskCreateOptions{CategoryID: &category.ID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteTask(task.ID) }) + loaded, err := d.GetTask(task.ID) + if err != nil || loaded == nil || loaded.CategoryID == nil || *loaded.CategoryID != category.ID || loaded.CategoryName != name { + t.Fatalf("categorized task=%+v err=%v", loaded, err) + } + + renamed, err := d.RenameTaskCategory(category.ID, name+" renamed") + if err != nil { + t.Fatal(err) + } + if renamed.TaskCount != 1 { + t.Fatalf("renamed task count=%d, want 1", renamed.TaskCount) + } + loaded, _ = d.GetTask(task.ID) + if loaded.CategoryName != renamed.Name { + t.Fatalf("task category name=%q, want %q", loaded.CategoryName, renamed.Name) + } + + if _, err := d.SetTaskCategory(task.ID, nil); err != nil { + t.Fatal(err) + } + loaded, _ = d.GetTask(task.ID) + if loaded.CategoryID != nil || loaded.CategoryName != "" { + t.Fatalf("uncategorized task retained category: %+v", loaded) + } + if _, err := d.SetTaskCategory(task.ID, &category.ID); err != nil { + t.Fatal(err) + } + deleted, err := d.DeleteTaskCategory(category.ID) + if err != nil || !deleted { + t.Fatalf("delete category=%v err=%v", deleted, err) + } + loaded, _ = d.GetTask(task.ID) + if loaded.CategoryID != nil || loaded.CategoryName != "" { + t.Fatalf("category deletion did not clear task: %+v", loaded) + } +} + +func TestCreateTaskRejectsMissingCategoryAtomically(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + defer d.Close() + + missing := int64(1 << 62) + description := fmt.Sprintf("missing-category-%d", time.Now().UnixNano()) + if _, err := d.CreateTaskWithOptions(description, "goal", TaskCreateOptions{CategoryID: &missing}); !errors.Is(err, ErrTaskCategoryNotFound) { + t.Fatalf("create error=%v, want %v", err, ErrTaskCategoryNotFound) + } + var count int + if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, description).Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("failed create leaked %d exploration rows", count) + } +} diff --git a/db/task_scope.go b/db/task_scope.go index 40e3a28..f584b5d 100644 --- a/db/task_scope.go +++ b/db/task_scope.go @@ -1,6 +1,7 @@ package db import ( + "database/sql" "fmt" "net" "strconv" @@ -10,17 +11,18 @@ import ( // TaskScope is one row of a task's test scope — the coverage denominator and the // per-task authorization edge. Rows come either from insertAssets (source='auto', // conservative, one per explicitly-inserted asset) or from the add_task_scope tool -// (source='agent', for company / whole-root-domain / subdomain / ip). +// (source='agent', for company / domain / network / ICP / keyword scope). type TaskScope struct { ID int64 `json:"id"` TaskID int64 `json:"task_id"` - Kind string `json:"kind"` // company|root_domain|subdomain|ip|cidr + Kind string `json:"kind"` // company|root_domain|subdomain|ip|cidr|icp|keyword CompanyID *int64 `json:"company_id,omitempty"` // CompanyName is resolved for kind=company so callers can label a scope row // without a second lookup. Empty when the row is not a company reference. CompanyName string `json:"company_name,omitempty"` Domain string `json:"domain,omitempty"` Net string `json:"net,omitempty"` + Value string `json:"value,omitempty"` Source string `json:"source"` Reason string `json:"reason,omitempty"` } @@ -53,11 +55,11 @@ func ipToHostCIDR(ip string) string { // upsertTaskScope inserts one scope row idempotently (uq_task_scope). A duplicate is // silently ignored. taskID<=0 or empty kind → no-op. -func (s *AssetStore) upsertTaskScope(ts TaskScope) error { +func (s *AssetStore) upsertTaskScopeResult(ts TaskScope) (bool, error) { if ts.TaskID <= 0 || ts.Kind == "" { - return nil + return false, nil } - var domainVal, netVal, companyVal any + var domainVal, netVal, companyVal, valueVal any if ts.Domain != "" { domainVal = ts.Domain } @@ -67,14 +69,35 @@ func (s *AssetStore) upsertTaskScope(ts TaskScope) error { if ts.CompanyID != nil && *ts.CompanyID > 0 { companyVal = *ts.CompanyID } + if ts.Value != "" { + valueVal = ts.Value + } src := ts.Source if src == "" { src = "auto" } - _, err := s.db.Exec(` -INSERT INTO task_scope(task_id, kind, company_id, domain, net, source, reason) -VALUES ($1,$2,$3,$4,$5::cidr,$6,NULLIF($7,'')) -ON CONFLICT DO NOTHING`, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, src, ts.Reason) + query := ` +INSERT INTO task_scope(task_id, kind, company_id, domain, net, value, source, reason) +VALUES ($1,$2,$3,$4,$5::cidr,$6,$7,NULLIF($8,'')) +ON CONFLICT DO NOTHING` + var ( + result sql.Result + err error + ) + if s.tx != nil { + result, err = s.tx.Exec(query, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, valueVal, src, ts.Reason) + } else { + result, err = s.db.Exec(query, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, valueVal, src, ts.Reason) + } + if err != nil { + return false, err + } + rows, err := result.RowsAffected() + return rows > 0, err +} + +func (s *AssetStore) upsertTaskScope(ts TaskScope) error { + _, err := s.upsertTaskScopeResult(ts) return err } @@ -182,8 +205,14 @@ func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason, source str return ts, fmt.Errorf("无效 ip/cidr: %s", value) } ts.Net = v + case "icp", "keyword": + parsed, err := ParseScopeInput(ScopeInput{Kind: kind, Value: value}) + if err != nil { + return ts, err + } + ts.Value = parsed.Value default: - return ts, fmt.Errorf("不支持的 kind: %s（company/root_domain/subdomain/ip/cidr）", kind) + return ts, fmt.Errorf("不支持的 kind: %s（company/root_domain/subdomain/ip/cidr/icp/keyword）", kind) } if err := s.upsertTaskScope(ts); err != nil { return ts, err @@ -206,7 +235,7 @@ func (s *AssetStore) DeleteTaskScope(taskID, scopeID int64) (bool, error) { func (s *AssetStore) ListTaskScope(taskID int64) ([]TaskScope, error) { rows, err := s.db.Query(` SELECT ts.id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''), - COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'') + COALESCE(ts.net::text,''), COALESCE(ts.value,''), ts.source, COALESCE(ts.reason,'') FROM task_scope ts LEFT JOIN companies c ON c.id=ts.company_id WHERE ts.task_id=$1 ORDER BY ts.id`, taskID) @@ -218,7 +247,7 @@ WHERE ts.task_id=$1 ORDER BY ts.id`, taskID) for rows.Next() { var t TaskScope var cid int64 - if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.CompanyName, &t.Domain, &t.Net, &t.Source, &t.Reason); err != nil { + if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.CompanyName, &t.Domain, &t.Net, &t.Value, &t.Source, &t.Reason); err != nil { return nil, err } t.TaskID = taskID @@ -279,9 +308,13 @@ target AS ( JOIN task_scope ts ON ts.task_id = $1 AND ( (ts.kind='company' AND a.company_id = ts.company_id) OR (ts.kind='root_domain' AND a.root_domain = ts.domain) - OR (ts.kind='subdomain' AND a.domain = ts.domain) - OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) - ) + OR (ts.kind='subdomain' AND a.domain = ts.domain) + OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet) + OR (ts.kind='icp' AND ( + lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value + OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value + )) + ) ), tested AS ( SELECT DISTINCT ea.asset_id diff --git a/db/tasks.go b/db/tasks.go index c7a94ac..bc0a57b 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -13,6 +13,8 @@ import ( type Task struct { ID int64 `json:"id"` Name string `json:"name"` // 可选任务名称;空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` + CategoryName string `json:"category_name,omitempty"` Description string `json:"description"` Goal string `json:"goal"` ExplorationID int64 `json:"exploration_id"` @@ -139,6 +141,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe // with the task/exploration row. type TaskCreateOptions struct { Name string // 可选任务名称;空=未命名 + CategoryID *int64 SourceTaskIDs []int64 CompanyIDs []int64 LLMProfileIDs []int64 @@ -189,13 +192,25 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); } opts.PlanHeartbeatSeconds = normalizeHeartbeat(opts.PlanHeartbeatSeconds) coverageEnabled := opts.CoverageEnabled == nil || *opts.CoverageEnabled + var categoryName string + if opts.CategoryID != nil { + if *opts.CategoryID <= 0 { + return nil, fmt.Errorf("%w: category id must be positive", ErrTaskCategoryInvalid) + } + if err := tx.QueryRow(`SELECT name FROM task_categories WHERE id=$1`, *opts.CategoryID).Scan(&categoryName); err != nil { + if err == sql.ErrNoRows { + return nil, ErrTaskCategoryNotFound + } + return nil, err + } + } var active *int64 if len(opts.LLMProfileIDs) > 0 { id := opts.LLMProfileIDs[0] active = &id } t := &Task{ - Name: opts.Name, + Name: opts.Name, CategoryID: opts.CategoryID, CategoryName: categoryName, Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: active, ActiveLLMProfileID: active, LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...), @@ -205,9 +220,9 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); CoverageEnabled: coverageEnabled, } if err := tx.QueryRow(` -INSERT INTO tasks(name, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled) -VALUES ($1,$2,$3,$4,$5,$5,$6,$7,$8) -RETURNING id, status, paused, created_at`, opts.Name, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, coverage_enabled) +VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9) +RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, coverageEnabled).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { @@ -287,11 +302,13 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { return nil } -const taskCols = `id, COALESCE(name,''), description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(coverage_enabled,true), first_run_at, deadline_at` +const taskCols = `id, COALESCE(name,''), category_id, +COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_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), COALESCE(coverage_enabled,true), first_run_at, deadline_at` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Name, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.CoverageEnabled, &t.FirstRunAt, &t.DeadlineAt); err != nil { return nil, err } return &t, nil diff --git a/server/assets.go b/server/assets.go index 5c31820..dc854b2 100644 --- a/server/assets.go +++ b/server/assets.go @@ -531,6 +531,9 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) { errs = append(errs, errEntry{Index: i, Error: err.Error()}) continue } + if req.TaskID > 0 { + _ = as.SetTaskAssetSource(req.TaskID, id, "api", "通过资产 API 登记", nil) + } results = append(results, result{Index: i, ID: id, Type: a.Type}) } writeJSON(w, 200, map[string]any{"results": results, "errors": errs}) diff --git a/server/dto.go b/server/dto.go index f30881c..91a3558 100644 --- a/server/dto.go +++ b/server/dto.go @@ -33,6 +33,8 @@ type TaskDTO struct { ID string `json:"id"` ExplorationID int64 `json:"exploration_id"` Name string `json:"name"` // 可选任务名称;空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` + CategoryName string `json:"category_name,omitempty"` Description string `json:"description"` Goal string `json:"goal"` Status string `json:"status"` // created | running | paused | done | failed @@ -85,6 +87,8 @@ func taskDTO(t *Task, status string) TaskDTO { ID: t.ID, ExplorationID: t.ExpID, Name: t.Name, + CategoryID: lifecycle.CategoryID, + CategoryName: lifecycle.CategoryName, Description: t.Description, Goal: t.Goal, Status: status, diff --git a/server/manager.go b/server/manager.go index 24bde8f..46cf066 100644 --- a/server/manager.go +++ b/server/manager.go @@ -27,15 +27,17 @@ 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"` - Name string `json:"name"` // 可选任务名称;空=未命名 - 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"` + Name string `json:"name"` // 可选任务名称;空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` + CategoryName string `json:"category_name,omitempty"` + 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. @@ -85,6 +87,8 @@ type taskLifecycleState struct { DeadlineAt int64 SourceTaskIDs []int64 CompanyIDs []int64 + CategoryID *int64 + CategoryName string } func (t *Task) lifecycleSnapshot() taskLifecycleState { @@ -108,6 +112,8 @@ func (t *Task) lifecycleSnapshotLocked() taskLifecycleState { DeadlineAt: t.DeadlineAt, SourceTaskIDs: append([]int64(nil), t.SourceTaskIDs...), CompanyIDs: append([]int64(nil), t.CompanyIDs...), + CategoryID: cloneInt64Ptr(t.CategoryID), + CategoryName: t.CategoryName, } } @@ -128,6 +134,8 @@ func (t *Task) updateLifecycle(update func(*taskLifecycleState)) { t.DeadlineAt = state.DeadlineAt t.SourceTaskIDs = append(t.SourceTaskIDs[:0], state.SourceTaskIDs...) t.CompanyIDs = append(t.CompanyIDs[:0], state.CompanyIDs...) + t.CategoryID = cloneInt64Ptr(state.CategoryID) + t.CategoryName = state.CategoryName t.lifecycleMu.Unlock() } @@ -688,7 +696,8 @@ func unixNanoOrZero(t *time.Time) int64 { func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task { return &Task{ ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID, - Name: pt.Name, + Name: pt.Name, + CategoryID: cloneInt64Ptr(pt.CategoryID), CategoryName: pt.CategoryName, 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, @@ -734,6 +743,75 @@ func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.Task return t, nil } +// RenameTaskCategory persists a category name and refreshes every live task DTO +// that references it. taskStateMu keeps this ordered with task reassignment. +func (m *Manager) RenameTaskCategory(id int64, name string) (*pgdb.TaskCategory, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + category, err := m.pg.RenameTaskCategory(id, name) + if err != nil { + return nil, err + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, task := range m.tasks { + task.updateLifecycle(func(state *taskLifecycleState) { + if state.CategoryID != nil && *state.CategoryID == id { + state.CategoryName = category.Name + } + }) + } + return category, nil +} + +// DeleteTaskCategory moves all affected live tasks to the uncategorized bucket. +func (m *Manager) DeleteTaskCategory(id int64) (bool, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + deleted, err := m.pg.DeleteTaskCategory(id) + if err != nil || !deleted { + return deleted, err + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, task := range m.tasks { + task.updateLifecycle(func(state *taskLifecycleState) { + if state.CategoryID != nil && *state.CategoryID == id { + state.CategoryID = nil + state.CategoryName = "" + } + }) + } + return true, nil +} + +// SetTaskCategory updates one live task without interrupting its runtime. +func (m *Manager) SetTaskCategory(taskID string, categoryID *int64) (*pgdb.TaskCategory, error) { + m.taskStateMu.Lock() + defer m.taskStateMu.Unlock() + id, err := strconv.ParseInt(taskID, 10, 64) + if err != nil || id <= 0 { + return nil, pgdb.ErrTaskCategoryTaskNotFound + } + category, err := m.pg.SetTaskCategory(id, categoryID) + if err != nil { + return nil, err + } + m.mu.RLock() + task := m.tasks[taskID] + m.mu.RUnlock() + if task != nil { + task.updateLifecycle(func(state *taskLifecycleState) { + state.CategoryID = cloneInt64Ptr(categoryID) + state.CategoryName = "" + if category != nil { + state.CategoryName = category.Name + } + }) + } + return category, 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 diff --git a/server/server.go b/server/server.go index 4ee4864..219ab50 100644 --- a/server/server.go +++ b/server/server.go @@ -390,9 +390,9 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开;每轮读) pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg) pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner")) - pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work - pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction - pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy + pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work + pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction + pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy pl.SetWebSearch(s.webSearchFor("planner")) pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开;每轮读) return pl, wk @@ -646,15 +646,23 @@ 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-categories", s.pgListTaskCategories) + mux.HandleFunc("POST /api/task-categories", s.pgCreateTaskCategory) + mux.HandleFunc("PATCH /api/task-categories/{id}", s.pgRenameTaskCategory) + mux.HandleFunc("DELETE /api/task-categories/{id}", s.pgDeleteTaskCategory) 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("PATCH /api/tasks/{id}/category", s.updateTaskCategory) 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) + mux.HandleFunc("POST /api/tasks/{id}/assets", s.attachTaskAssets) + mux.HandleFunc("DELETE /api/tasks/{id}/assets/{assetID}", s.detachTaskAsset) + mux.HandleFunc("GET /api/tasks/{id}/intent-assets", s.taskIntentAssets) // 工作空间文件管理器（针对 workDir） mux.HandleFunc("GET /api/workspace/list", s.wsList) @@ -667,10 +675,10 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/tasks/{id}/scope", s.taskScopeList) mux.HandleFunc("POST /api/tasks/{id}/scope", s.taskScopeAdd) mux.HandleFunc("DELETE /api/tasks/{id}/scope/{sid}", s.taskScopeDelete) - mux.HandleFunc("GET /api/tasks/{id}/goals", s.listGoals) // 目标管理:列出本任务全部目标 - mux.HandleFunc("POST /api/tasks/{id}/goals", s.addGoal) // 目标管理:人工新增目标(复活任务) - mux.HandleFunc("PATCH /api/tasks/{id}/goals/{gid}", s.editGoal) // 目标管理:修改目标(复活任务) - mux.HandleFunc("DELETE /api/tasks/{id}/goals/{gid}", s.deleteGoal) // 目标管理:硬删除目标(不复活) + mux.HandleFunc("GET /api/tasks/{id}/goals", s.listGoals) // 目标管理:列出本任务全部目标 + mux.HandleFunc("POST /api/tasks/{id}/goals", s.addGoal) // 目标管理:人工新增目标(复活任务) + mux.HandleFunc("PATCH /api/tasks/{id}/goals/{gid}", s.editGoal) // 目标管理:修改目标(复活任务) + mux.HandleFunc("DELETE /api/tasks/{id}/goals/{gid}", s.deleteGoal) // 目标管理:硬删除目标(不复活) mux.HandleFunc("GET /api/tasks/{id}/constraints", s.listConstraints) // 约束管理:列出本任务操作约束 mux.HandleFunc("POST /api/tasks/{id}/constraints", s.addConstraint) // 约束管理:新增约束(不通知 planner) mux.HandleFunc("PATCH /api/tasks/{id}/constraints/{cid}", s.editConstraint) // 约束管理:修改约束 @@ -1315,6 +1323,7 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { type createTaskReq struct { Name string `json:"name,omitempty"` // 可选任务名称;省略/空=未命名 + CategoryID *int64 `json:"category_id,omitempty"` Description string `json:"description"` Goal string `json:"goal"` LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置 @@ -1372,12 +1381,16 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { } req.CompanyIDs = companyIDs t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{ - Name: strings.TrimSpace(req.Name), + Name: strings.TrimSpace(req.Name), CategoryID: req.CategoryID, SourceTaskIDs: sourceIDs, CompanyIDs: req.CompanyIDs, LLMProfileIDs: req.LLMProfileIDs, TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds, CoverageEnabled: req.CoverageEnabled, }) if err != nil { + if errors.Is(err, db.ErrTaskCategoryInvalid) || errors.Is(err, db.ErrTaskCategoryNotFound) { + writeErr(w, 400, "任务分类不存在或无效") + return + } if errors.Is(err, db.ErrTaskCompanyIDsInvalid) || errors.Is(err, db.ErrTaskCompanyNotFound) { writeErr(w, 400, "关联企业不存在或无效") return @@ -1554,6 +1567,9 @@ func (s *Server) seed(t *Task, text string) { } else { rootID, _ = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: taskID}) } + if rootID > 0 { + _ = as.SetTaskAssetSource(taskID, rootID, "task", "由任务描述或目标初始化", nil) + } } // anchor the seeded assets to this task's begin root as lineage/provenance // (the asset graph is global and shared; anchoring no longer gates reads). diff --git a/server/task_assets.go b/server/task_assets.go new file mode 100644 index 0000000..cda4287 --- /dev/null +++ b/server/task_assets.go @@ -0,0 +1,108 @@ +package server + +import ( + "errors" + "net/http" + "strings" + + "github.com/Autumn-27/artex/db" +) + +const maxTaskAssetRequestBytes = 512 << 10 + +func writeTaskAssetError(w http.ResponseWriter, err error) { + switch { + case errors.Is(err, db.ErrTaskAssetInvalid): + writeErr(w, http.StatusBadRequest, err.Error()) + case errors.Is(err, db.ErrTaskAssetTaskNotFound), errors.Is(err, db.ErrTaskAssetAssetNotFound): + writeErr(w, http.StatusNotFound, err.Error()) + default: + writeErr(w, http.StatusInternalServerError, err.Error()) + } +} + +func (s *Server) attachTaskAssets(w http.ResponseWriter, r *http.Request) { + task, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, http.StatusNotFound, "task not found") + return + } + r.Body = http.MaxBytesReader(w, r.Body, maxTaskAssetRequestBytes) + var request struct { + AssetIDs []int64 `json:"asset_ids"` + SourceSummary string `json:"source_summary"` + Scope companyScopeInputs `json:"scope"` + } + if err := decode(r, &request); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大") + } else { + writeErr(w, http.StatusBadRequest, err.Error()) + } + return + } + request.SourceSummary = strings.TrimSpace(request.SourceSummary) + taskID, _ := parseTaskID(task.ID) + if request.Scope != nil && len(request.AssetIDs) > 0 { + writeErr(w, http.StatusBadRequest, "scope 与 asset_ids 不能同时提交") + return + } + if request.Scope != nil { + mutation, err := s.m.Assets().RegisterTaskAssetScopes(taskID, request.Scope) + if err != nil { + writeTaskAssetError(w, err) + return + } + task.Notify() + writeJSON(w, http.StatusOK, mutation) + return + } + mutation, err := s.m.Assets().AttachAssetsToTask(taskID, request.AssetIDs, request.SourceSummary) + if err != nil { + writeTaskAssetError(w, err) + return + } + task.Notify() + writeJSON(w, http.StatusOK, mutation) +} + +func (s *Server) detachTaskAsset(w http.ResponseWriter, r *http.Request) { + task, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, http.StatusNotFound, "task not found") + return + } + assetID, ok := pathInt(r, "assetID") + if !ok { + writeErr(w, http.StatusBadRequest, "bad asset id") + return + } + taskID, _ := parseTaskID(task.ID) + detached, err := s.m.Assets().DetachAssetFromTask(taskID, assetID) + if err != nil { + writeTaskAssetError(w, err) + return + } + if !detached { + writeErr(w, http.StatusNotFound, "asset is not associated with this task") + return + } + task.Notify() + writeJSON(w, http.StatusOK, map[string]any{"detached": assetID}) +} + +func (s *Server) taskIntentAssets(w http.ResponseWriter, r *http.Request) { + task, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, http.StatusNotFound, "task not found") + return + } + taskID, _ := parseTaskID(task.ID) + assets, err := s.m.Assets().IntentAssets(taskID) + if err != nil { + writeTaskAssetError(w, err) + return + } + writeJSON(w, http.StatusOK, map[string]any{"assets": assets}) +} diff --git a/server/task_categories.go b/server/task_categories.go new file mode 100644 index 0000000..9cca440 --- /dev/null +++ b/server/task_categories.go @@ -0,0 +1,159 @@ +package server + +import ( + "bytes" + "encoding/json" + "errors" + "net/http" + "strings" + "unicode/utf8" + + "github.com/Autumn-27/artex/db" +) + +const maxTaskCategoryRequestBytes = 16 << 10 + +type taskCategoryRequest struct { + Name *string `json:"name"` +} + +func decodeTaskCategoryRequest(w http.ResponseWriter, r *http.Request) (*taskCategoryRequest, bool) { + r.Body = http.MaxBytesReader(w, r.Body, maxTaskCategoryRequestBytes) + var request taskCategoryRequest + if err := decode(r, &request); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大") + } else { + writeErr(w, http.StatusBadRequest, err.Error()) + } + return nil, false + } + if request.Name == nil || strings.TrimSpace(*request.Name) == "" { + writeErr(w, http.StatusBadRequest, "分类名称不能为空") + return nil, false + } + if utf8.RuneCountInString(strings.TrimSpace(*request.Name)) > db.MaxTaskCategoryNameRunes { + writeErr(w, http.StatusBadRequest, "分类名称最多 80 个字符") + return nil, false + } + return &request, true +} + +func writeTaskCategoryError(w http.ResponseWriter, err error) { + switch { + case errors.Is(err, db.ErrTaskCategoryInvalid): + writeErr(w, http.StatusBadRequest, err.Error()) + case errors.Is(err, db.ErrTaskCategoryNameConflict): + writeErr(w, http.StatusConflict, "分类名称已存在") + case errors.Is(err, db.ErrTaskCategoryNotFound), errors.Is(err, db.ErrTaskCategoryTaskNotFound): + writeErr(w, http.StatusNotFound, err.Error()) + default: + writeErr(w, http.StatusInternalServerError, err.Error()) + } +} + +func (s *Server) pgListTaskCategories(w http.ResponseWriter, _ *http.Request) { + pg := s.pg(w) + if pg == nil { + return + } + categories, err := pg.ListTaskCategories() + if err != nil { + writeErr(w, http.StatusInternalServerError, err.Error()) + return + } + writeJSON(w, http.StatusOK, map[string]any{"categories": categories}) +} + +func (s *Server) pgCreateTaskCategory(w http.ResponseWriter, r *http.Request) { + pg := s.pg(w) + if pg == nil { + return + } + request, ok := decodeTaskCategoryRequest(w, r) + if !ok { + return + } + category, err := pg.CreateTaskCategory(*request.Name) + if err != nil { + writeTaskCategoryError(w, err) + return + } + writeJSON(w, http.StatusCreated, category) +} + +func (s *Server) pgRenameTaskCategory(w http.ResponseWriter, r *http.Request) { + id, ok := pathInt(r, "id") + if !ok { + writeErr(w, http.StatusBadRequest, "bad task category id") + return + } + request, ok := decodeTaskCategoryRequest(w, r) + if !ok { + return + } + category, err := s.m.RenameTaskCategory(id, *request.Name) + if err != nil { + writeTaskCategoryError(w, err) + return + } + writeJSON(w, http.StatusOK, category) +} + +func (s *Server) pgDeleteTaskCategory(w http.ResponseWriter, r *http.Request) { + id, ok := pathInt(r, "id") + if !ok { + writeErr(w, http.StatusBadRequest, "bad task category id") + return + } + deleted, err := s.m.DeleteTaskCategory(id) + if err != nil { + writeTaskCategoryError(w, err) + return + } + if !deleted { + writeErr(w, http.StatusNotFound, "task category not found") + return + } + writeJSON(w, http.StatusOK, map[string]any{"deleted": id}) +} + +func (s *Server) updateTaskCategory(w http.ResponseWriter, r *http.Request) { + taskID := r.PathValue("id") + if _, ok := s.m.Task(taskID); !ok { + writeErr(w, http.StatusNotFound, "task not found") + return + } + r.Body = http.MaxBytesReader(w, r.Body, maxTaskCategoryRequestBytes) + var request struct { + CategoryID json.RawMessage `json:"category_id"` + } + if err := decode(r, &request); err != nil { + writeErr(w, http.StatusBadRequest, err.Error()) + return + } + if len(request.CategoryID) == 0 { + writeErr(w, http.StatusBadRequest, "category_id is required") + return + } + var categoryID *int64 + if !bytes.Equal(bytes.TrimSpace(request.CategoryID), []byte("null")) { + var id int64 + if err := json.Unmarshal(request.CategoryID, &id); err != nil { + writeErr(w, http.StatusBadRequest, "任务分类 id 无效") + return + } + categoryID = &id + } + if categoryID != nil && *categoryID <= 0 { + writeErr(w, http.StatusBadRequest, "任务分类 id 无效") + return + } + if _, err := s.m.SetTaskCategory(taskID, categoryID); err != nil { + writeTaskCategoryError(w, err) + return + } + task, _ := s.m.Task(taskID) + writeJSON(w, http.StatusOK, taskDTO(task, s.resolvedTaskStatus(task))) +} diff --git a/web/src/app/(main)/function/assets/page.tsx b/web/src/app/(main)/function/assets/page.tsx index e00026f..b9b1b76 100644 --- a/web/src/app/(main)/function/assets/page.tsx +++ b/web/src/app/(main)/function/assets/page.tsx @@ -19,6 +19,7 @@ import { } from "lucide-react"; import { toast } from "sonner"; +import { ScopeTextEditor } from "@/components/scope-text-editor"; import { AlertDialog, AlertDialogAction, @@ -43,7 +44,7 @@ import { DialogTitle, DialogTrigger, } from "@/components/ui/dialog"; -import { Field, FieldDescription, FieldError, FieldGroup, FieldLabel } from "@/components/ui/field"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { @@ -57,10 +58,9 @@ import { } from "@/components/ui/sheet"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { Textarea } from "@/components/ui/textarea"; import { api } from "@/lib/api"; -import { MAX_COMPANY_SCOPE_RULES, type ParsedCompanyScopeText, parseCompanyScopeText } from "@/lib/company-scope"; -import type { Asset, Company, CompanyScopeKind, CompanyScopeRule } from "@/lib/types"; +import { MAX_COMPANY_SCOPE_RULES, parseCompanyScopeText } from "@/lib/company-scope"; +import type { Asset, Company, CompanyScopeRule } from "@/lib/types"; import { cn } from "@/lib/utils"; // ── DSL autocomplete ────────────────────────────────────────────────────────── @@ -1321,70 +1321,6 @@ function CompanyAvatar({ name, logo }: { name: string; logo?: string }) { ); } -const COMPANY_SCOPE_KIND_LABELS: Record = { - domain: "域名", - ip: "IP", - cidr: "CIDR", - icp: "ICP", - keyword: "关键词", -}; - -function ScopeTextEditor({ - id, - value, - onValueChange, - parsed, -}: { - id: string; - value: string; - onValueChange: (value: string) => void; - parsed: ParsedCompanyScopeText; -}) { - const counts = React.useMemo(() => { - const result = new Map(); - for (const rule of parsed.rules) result.set(rule.kind, (result.get(rule.kind) ?? 0) + 1); - return result; - }, [parsed.rules]); - return ( - 0}> - 资产范围 - 每行一条，自动识别域名、IP、CIDR、ICP 备案和企业关键词。 -