From b7e0b930237a443202e8633c9234768ccbe1c33f Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 21 Aug 2026 11:16:19 +0800 Subject: [PATCH 1/8] init --- CHANGELOG.md | 14 +- agent/chat.go | 1 + agent/mainagent.go | 1 + agent/planner.go | 1 + agent/provider.go | 26 +- agent/runinfo.go | 47 ++ agent/tools_insert.go | 2 +- agent/worker.go | 13 + db/commands.go | 12 + db/config.go | 28 +- db/db.go | 4 +- db/exploration.go | 10 +- db/exploration_sources_test.go | 6 +- db/findings.go | 51 ++ db/llm_usage.go | 183 ++++++ db/schema.sql | 47 +- db/skill_usage.go | 200 ++++++ db/skill_usage_test.go | 116 ++++ db/task_scope.go | 21 +- db/tasks.go | 2 + llmrec/llmrec.go | 134 +++- report/findings.go | 165 +++++ server/assembly.go | 6 +- server/commands.go | 49 ++ server/dto.go | 2 + server/engine.go | 3 +- server/manager.go | 3 + server/server.go | 203 +++++- server/server_mgmt.go | 59 +- server/skill_usage.go | 88 +++ web/src/app/(main)/chat/page.tsx | 17 +- web/src/app/(main)/dashboard/page.tsx | 594 +++++++++--------- .../(main)/function/findings/detail/page.tsx | 6 +- web/src/app/(main)/function/findings/page.tsx | 204 +++++- .../tasks/detail/_tabs/overview-tab.tsx | 388 +++++++++--- .../tasks/detail/_tabs/report-tab.tsx | 15 +- .../tasks/detail/_tabs/sessions-tab.tsx | 69 +- web/src/app/(main)/function/tasks/page.tsx | 438 ++++++++++--- web/src/app/(main)/system/llm/page.tsx | 100 +-- web/src/app/(main)/system/skills/page.tsx | 487 ++++++++++++-- web/src/components/copy-button.tsx | 67 ++ web/src/lib/api.ts | 358 ++++++++--- web/src/lib/mock/data.ts | 24 +- web/src/lib/mock/handler.ts | 121 +++- web/src/lib/types.ts | 117 +++- web/src/lib/utils.ts | 29 + 46 files changed, 3632 insertions(+), 899 deletions(-) create mode 100644 agent/runinfo.go create mode 100644 db/llm_usage.go create mode 100644 db/skill_usage.go create mode 100644 db/skill_usage_test.go create mode 100644 report/findings.go create mode 100644 server/skill_usage.go create mode 100644 web/src/components/copy-button.tsx diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c53e7e..40928b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ 本项目的重要变更记录在此文件中，格式参考 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/)。 -## [Unreleased] - 2026-08-20 +## [0.3.2] - 2026-08-20 ### Added @@ -23,9 +23,13 @@ - 任务详情页按真实 LLM 调用状态展示运行、空闲和暂停状态；来源任务的事实、漏洞、意图、资产引用和图谱节点统一标记来源并保持只读。 - Agent 预览、任务报告和相关详情视图统一使用共享 Markdown 渲染组件。 - LLM 模型配置采用卡片和抽屉交互，模型列表在抽屉内支持独立滚动。 +- 任务暂停不再拦截主 Agent 对话：主 Agent 编排会话独立于任务暂停，暂停中仍可继续发送新消息（暂停只终止当前进行中的那一轮）。 +- Worker 单次 run 的墙钟默认时长由 600 秒调整为 1200 秒。 ### Fixed +- 移除主 Agent 控制台的乐观回显，改为纯服务端数据渲染，修复暂停或发送失败等场景下消息错乱、串入其他会话内容的问题。 +- 打开主 Agent 会话默认滚动到底并完整显示最后一条回复：最后一条回复的完整内容懒加载展开后自动贴底，不再被顶出屏幕。 - 修复暂停任务后主 Agent 会话仍继续运行，以及编排 Agent 暂停任务时未同步停止主 Agent 的问题。 - 修复任务完成、Planner、Worker 或主 Agent 实际运行时，任务状态徽标与操作按钮状态不同步的问题。 - 修复 Worker 取消、任务删除和并发写入之间可能产生迟到黑板写入或残留文件的问题。 @@ -34,4 +38,10 @@ - 修复后端关闭时父 context 可能抢先覆盖具名 `shutdown` 原因的竞态，以及中文活动摘要按字节截断导致乱码的问题。 - 修复带部分流式输出的取消事件丢失真实终止原因的问题，同时保留取消前已生成的内容。 -[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.1...HEAD +### 贡献者 + +- [@Autumn-27](https://github.com/Autumn-27) +- [@neouks](https://github.com/neouks) + +[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.2...HEAD +[0.3.2]: https://github.com/Autumn-27/ARTEX/compare/v0.3.1...v0.3.2 diff --git a/agent/chat.go b/agent/chat.go index 86deccc..41be93b 100644 --- a/agent/chat.go +++ b/agent/chat.go @@ -92,6 +92,7 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin // visible skills/MCP and lets the DB tools table filter/override. DefaultTools // have no tools-table rows, so they always pass through. base := actool.DefaultTools() + ctx = WithRunInfo(ctx, RunInfo{SessionID: sessionID}) tools, def, cleanup := AugmentTools(ctx, agentKey, base) defer cleanup() diff --git a/agent/mainagent.go b/agent/mainagent.go index 86b198b..b65d8c0 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -85,6 +85,7 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t tsx.SetNotifyGoal(notifyGoal) // set_goals 新增目标 → 给 planner 记一条「人新增了目标：…」触发 // 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash） base := append(tsx.MainAgentTools(), actool.DefaultTools()...) + ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)}) tools, def, cleanup := AugmentTools(ctx, "mainagent", base) defer cleanup() // 本任务的工作目录 /tasks/，先建好。 diff --git a/agent/planner.go b/agent/planner.go index b2757fe..d460e6c 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -273,6 +273,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts } // 领域工具 + 基础默认工具集（Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash） base := append(tsx.PlannerTools(), actool.DefaultTools()...) + ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)}) tools, def, cleanup := AugmentTools(ctx, "planner", base) defer cleanup() // 关键态势（刚完成的意图 + 预取的完整图）改放【本轮 user 输入】(见下方 input)，system diff --git a/agent/provider.go b/agent/provider.go index 465812b..6c1c8b5 100644 --- a/agent/provider.go +++ b/agent/provider.go @@ -41,9 +41,14 @@ type Config struct { // ContextWindowK is the model's context window in K tokens (user-configured), // used to size compaction thresholds. 0 = default; see CompactionWindow. ContextWindowK int - // ReasoningEffort selects the thinking mode: "" = 默认(不发送思考参数); - // "off" = 关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" = 开启并设强度. - // NewProvider derives the provider-specific request fields from it. + // ThinkingType 独立控制思考「开关」字段(thinking.type): + // "" = 不发送(默认,兼容不支持该字段的模型); "disabled" = 显式关闭; + // "enabled" = 开启. 与 ReasoningEffort 完全解耦——有些接口没有 thinking 字段、 + // 只靠强度参数就能激活思考,故两者可各自单独设置. + ThinkingType string + // ReasoningEffort 独立控制思考「强度」字段: + // "" = 不发送(默认); "low"/"medium"/"high"/"xhigh"/"max" = 对应强度. + // OpenAI 映射为顶层 reasoning_effort;Anthropic 映射为 output_config.effort. ReasoningEffort string } @@ -184,17 +189,10 @@ func (c Config) NewProvider() (llm.Provider, error) { Model: c.Model, HTTPClient: client, } - // Derive the provider thinking params from the single UI selector: - // "" → omit both (provider default); "off" → disable; else → enable + effort. - switch c.ReasoningEffort { - case "": - // leave both empty → omitted from the request - case "off": - lc.ThinkingType = "disabled" - default: - lc.ThinkingType = "enabled" - lc.ReasoningEffort = c.ReasoningEffort - } + // 思考开关与强度两个字段各自透传(空 = 该字段不发送)。二者解耦: + // 可只发 thinking.type、只发 effort、都发、或都不发。 + lc.ThinkingType = c.ThinkingType + lc.ReasoningEffort = c.ReasoningEffort if c.RatePerSecond > 0 || c.RatePerMinute > 0 { lc.RateLimit = &llm.RateLimit{PerSecond: c.RatePerSecond, PerMinute: c.RatePerMinute} } diff --git a/agent/runinfo.go b/agent/runinfo.go new file mode 100644 index 0000000..f006b6a --- /dev/null +++ b/agent/runinfo.go @@ -0,0 +1,47 @@ +package agent + +import ( + "context" + + "github.com/Autumn-27/artex/db" +) + +// RunInfo identifies WHICH run a tool call belongs to. Tool assembly only receives +// (ctx, agentKey) — the task/exploration ids live in the caller's arguments, not the +// ctx — so anything wired at assembly time (currently the Skill ledger in +// server/assembly.go) has no way to attribute a call to a task. Each run attaches +// its own RunInfo before calling AugmentTools; the wiring closure reads it once and +// captures it, so per-run attribution stays correct without threading parameters +// through the tool layer. Same pattern as TaskClock (see taskclock.go). +// +// Zero value = attribution unknown; every consumer must treat it as optional. +type RunInfo struct { + TaskID int64 // task registry id; 0 for non-task runs (chat sessions) + ExplorationID int64 // exploration id; 0 when unknown + IntentID int64 // worker's intent node; 0 for planner/mainagent/chat + SessionID string // chat conversation id; empty for task runs +} + +// explorationID reads a store's exploration id, tolerating a nil store (planner and +// worker runs can be driven without one in tests). +func explorationID(ts *db.ExplorationStore) int64 { + if ts == nil { + return 0 + } + return ts.ID() +} + +type runInfoKey struct{} + +// WithRunInfo attaches run attribution to ctx. +func WithRunInfo(ctx context.Context, ri RunInfo) context.Context { + return context.WithValue(ctx, runInfoKey{}, ri) +} + +// RunInfoFrom reads the RunInfo (zero value if none attached). +func RunInfoFrom(ctx context.Context) RunInfo { + if v, ok := ctx.Value(runInfoKey{}).(RunInfo); ok { + return v + } + return RunInfo{} +} diff --git a/agent/tools_insert.go b/agent/tools_insert.go index 65f7cf2..f2b943c 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -382,7 +382,7 @@ func (t *ToolSet) addTaskScope() actool.CoreTool { var added []map[string]any errs := map[string]string{} for i, e := range items { - ts, err := t.as.AddAgentScope(t.taskID, strings.TrimSpace(e.Kind), e.Value, a.Reason) + ts, err := t.as.AddAgentScope(t.taskID, strings.TrimSpace(e.Kind), e.Value, a.Reason, "agent") if err != nil { errs[strconv.Itoa(i)] = err.Error() continue diff --git a/agent/worker.go b/agent/worker.go index 47e25e9..258b05b 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -265,6 +265,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. // phase, Bash is hidden via Settlement.DisabledTools (no local gating needed). base := append(tsx.WorkerTools(), w.extraTools...) base = append(base, actool.DefaultTools()...) + ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts), IntentID: intent.ID}) tools, def, cleanup := AugmentTools(ctx, "worker", base) defer cleanup() @@ -360,6 +361,18 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. s := agentcore.NewSession(opts) defer s.Close() // release the session's background-task manager (temp dir + processes) + // Resume prior conversation if this intent was paused/blocked/exhausted and is + // being re-run. The transcript ID is deterministic per intent, so if a prior + // session exists the worker continues from where it left off instead of + // restarting from scratch. + if w.tx != nil { + _ = s.Resume(opts.SessionID) + if len(s.Messages()) > 0 { + seedUnlockFromHistory(s.Messages(), def.UnlockSkill) + input = "继续执行。" + } + } + // Budgets + settlement are owned by the SDK (MaxTurns/MaxDuration + Settlement): // on hit it runs a wrap-up turn and finishes with ReasonMaxTurns/ReasonTimeout. // We only add a HARD ctx backstop past the soft wall-clock budget, for the rare diff --git a/db/commands.go b/db/commands.go index 82b9c45..27172da 100644 --- a/db/commands.go +++ b/db/commands.go @@ -209,6 +209,18 @@ func (d *DB) ListLLMRecords(model, session, task string, page, size int) ([]LLMR return out, total, rows.Err() } +// ModelTokenStat is one model's aggregated token usage for a task, summed from the +// llm_usage metering ledger (see db/llm_usage.go). Calls is the number of LLM calls +// that hit this model. +type ModelTokenStat struct { + Model string `json:"model"` + Calls int `json:"calls"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheReadTokens int `json:"cache_read_tokens"` + CacheWriteTokens int `json:"cache_write_tokens"` +} + // LLMTask is one distinct task with its LLM-record count. type LLMTask struct { TaskID string `json:"task_id"` diff --git a/db/config.go b/db/config.go index 8d91487..40684d6 100644 --- a/db/config.go +++ b/db/config.go @@ -25,9 +25,11 @@ type LLMProfile struct { // ContextWindowK is the model's context window in K tokens, used to size // compaction thresholds. 0 = use a 200K default; capped at 1000 (1M). ContextWindowK int `json:"context_window_k"` - // ReasoningEffort is the thinking-mode selector. "" = 默认(不发送思考参数); - // "off" = 显式关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" = - // 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider. + // ThinkingType 独立控制思考「开关」(thinking.type):"" = 不发送(默认); + // "disabled" = 显式关闭; "enabled" = 开启. 与 ReasoningEffort 解耦. + ThinkingType string `json:"thinking_type"` + // ReasoningEffort 独立控制思考「强度」:"" = 不发送(默认); + // "low"/"medium"/"high"/"xhigh"/"max" = 对应强度. 见 agent.Config.NewProvider. ReasoningEffort string `json:"reasoning_effort"` IsDefault bool `json:"is_default"` // Priority orders the failover chain: higher goes first. The ACTIVE profile @@ -41,14 +43,14 @@ type LLMProfile struct { // profileCols is the read column list (hint variant, no api key) shared by the // list query; profileColsKey is the same with api_key for the single-row loads. -const profileCols = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude` -const profileColsKey = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude` +const profileCols = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude,COALESCE(thinking_type,'')` +const profileColsKey = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude,COALESCE(thinking_type,'')` // scanProfile reads one row in the profileCols / profileColsKey column order. The // 7th column lands in APIKeyHint or APIKey depending on which list the caller used. func scanProfile(sc interface{ Scan(...any) error }, into *string, p *LLMProfile) error { return sc.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, into, - &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault, &p.Priority, &p.PoolExclude) + &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault, &p.Priority, &p.PoolExclude, &p.ThinkingType) } func (d *DB) ListProfiles() ([]*LLMProfile, error) { @@ -121,18 +123,18 @@ func (d *DB) SaveProfile(p *LLMProfile) (int64, error) { } if p.ID == 0 { var id int64 - err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort,priority,pool_exclude) -VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13) RETURNING id`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude).Scan(&id) + err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort,priority,pool_exclude,thinking_type) +VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14) RETURNING id`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType).Scan(&id) return id, err } if p.APIKey == "" { - _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9,priority=$10,pool_exclude=$11 WHERE id=$12`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ID) + _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9,priority=$10,pool_exclude=$11,thinking_type=$12 WHERE id=$13`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.ID) return p.ID, err } - _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11,priority=$12,pool_exclude=$13 WHERE id=$14`, - p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ID) + _, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11,priority=$12,pool_exclude=$13,thinking_type=$14 WHERE id=$15`, + p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.ID) return p.ID, err } diff --git a/db/db.go b/db/db.go index 6a9e085..6a35129 100644 --- a/db/db.go +++ b/db/db.go @@ -103,7 +103,7 @@ type builtinAgent struct { key, name, role, desc string vars []promptVar interactiveShell bool // 建行时的默认交互式 shell 开关；ON CONFLICT 不覆盖用户后续手动开关 - runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用表默认(600)，0=不限时 + runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用种子默认(1200)，0=不限时 } type promptVar struct{ name, desc, example, source string } @@ -144,7 +144,7 @@ func (d *DB) seedBuiltins() error { var agentID int64 err := d.QueryRow(` INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds) -VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 600)) +VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 1200)) ON CONFLICT (key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID) if err != nil { diff --git a/db/exploration.go b/db/exploration.go index be01753..dbcdb35 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -215,7 +215,7 @@ func (s *ExplorationStore) SetNodeState(id int64, state string) error { // ResetRunningIntents returns any intent left in 'running' back to 'open' so it // is re-claimed. Called on startup: a 'running' intent with no live worker (a // backend restart or crashed worker goroutine left it stuck) would otherwise spin -// forever in the UI. Workers re-run an intent from scratch, so reopening is safe. +// forever in the UI. Workers resume from their transcript, so reopening is safe. func (s *ExplorationStore) ResetRunningIntents() (int64, error) { res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) @@ -227,9 +227,9 @@ WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID) } // ReopenIntent flips ONE not-successfully-finished intent (blocked/exhausted/stopped) -// back to 'open' so a worker re-claims and re-runs it from scratch (graph writes it -// already produced stay). done/open/running are left untouched. Returns whether a row -// changed (false = intent absent or not in a rerunnable state). Used by the "重跑" button. +// back to 'open' so a worker re-claims it (graph writes it already produced stay). +// The worker resumes from its prior LLM transcript instead of restarting from scratch. +// done/open/running are left untouched. Returns whether a row changed. Used by "重跑". func (s *ExplorationStore) ReopenIntent(id int64) (bool, error) { res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL @@ -244,7 +244,7 @@ WHERE id=$1 AND exploration_id=$2 AND kind='intent' // ReopenBlockedIntents flips EVERY 'blocked' intent in this exploration back to 'open' // (batch rerun after e.g. an LLM/network outage that blocked many at once). Returns the -// number reopened. Workers re-run each from scratch; kept graph writes remain. +// number reopened. Workers resume from their transcripts; kept graph writes remain. func (s *ExplorationStore) ReopenBlockedIntents() (int64, error) { res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL, blocked_reason=NULL WHERE exploration_id=$1 AND kind='intent' AND state='blocked'`, s.expID) diff --git a/db/exploration_sources_test.go b/db/exploration_sources_test.go index 7c66fae..07ceb29 100644 --- a/db/exploration_sources_test.go +++ b/db/exploration_sources_test.go @@ -356,13 +356,13 @@ func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) { _ = d.DeleteTask(grand.ID) _, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID}) }) - if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test"); err != nil { + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test", "agent"); err != nil { t.Fatal(err) } - if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test"); err != nil { + if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceUntestedDomain, "test", "agent"); err != nil { t.Fatal(err) } - if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test"); err != nil { + if _, err := assets.AddAgentScope(grand.ID, "root_domain", grandDomain, "test", "agent"); err != nil { t.Fatal(err) } if _, err := d.Exploration(source.ExplorationID).AddNode(KindFact, map[string]any{"summary": "tested source asset"}, 0, "confirmed", "worker", []int64{sourceTestedID}); err != nil { diff --git a/db/findings.go b/db/findings.go index 9dfe6e2..48f505c 100644 --- a/db/findings.go +++ b/db/findings.go @@ -210,6 +210,57 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding return out, total, err } +// ListFindingsForExport returns findings for the 发现 page 导出功能，携带完整 +// report 字段、不分页。ids 非空时按这批 finding id 精确导出(勾选导出),忽略 +// filter;ids 为空时按 filter 导出(导出当前筛选/全部)。结果按严重等级降序、 +// 再按时间倒序,与「导出汇总报告」的分组顺序一致。 +func (d *DB) ListFindingsForExport(f FindingFilter, ids []int64) ([]*DBFinding, error) { + const order = `ORDER BY 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` + cols := findingSelectCols + `, COALESCE(f.report, '')` + + var q string + var args []any + if len(ids) > 0 { + ph := make([]string, len(ids)) + for i, id := range ids { + ph[i] = fmt.Sprintf("$%d", i+1) + args = append(args, id) + } + q = `SELECT ` + cols + ` + FROM findings f + LEFT JOIN tasks t ON f.task_id = t.id + WHERE f.id IN (` + strings.Join(ph, ",") + `) + ` + order + } else { + where, wargs := f.where() + q = `SELECT ` + cols + ` + FROM findings f + LEFT JOIN tasks t ON f.task_id = t.id` + where + ` + ` + order + args = wargs + } + + rows, err := d.Query(q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + var out []*DBFinding + for rows.Next() { + f := &DBFinding{} + var aidsJSON string + if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity, + &f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, + &f.TaskDescription, &f.Report); err != nil { + return nil, err + } + _ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs) + out = append(out, f) + } + return out, rows.Err() +} + // FindingStats is the whole-table aggregate powering the 发现 page's stat cards // and vuln-class filter — computed server-side so it stays exact regardless of // pagination. diff --git a/db/llm_usage.go b/db/llm_usage.go new file mode 100644 index 0000000..ecbf83a --- /dev/null +++ b/db/llm_usage.go @@ -0,0 +1,183 @@ +package db + +import "strconv" + +// LLMUsage is one lightweight LLM-call metering row — the always-on usage ledger, +// distinct from llm_records (which stores full request/response bodies and is a +// gated debug feature). One row per completion call, written on both success and +// error, so token accounting is complete even for interrupted/failed runs. Carries +// only the dimensions needed to slice token spend (model / profile / task / agent), +// never any prompt or response content. +type LLMUsage struct { + TaskID string `json:"task_id"` // task registry id (matches llm_records.task_id) + ExplorationID int64 `json:"exploration_id"` // exploration id parsed from the session (0 = unknown/non-task) + Worker string `json:"worker"` // agent lane: worker / planner / mainagent / goals + Model string `json:"model"` + ProfileName string `json:"profile_name"` + LatencyMs int `json:"latency_ms"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheRead int `json:"cache_read"` + CacheWrite int `json:"cache_write"` + Status string `json:"status"` // ok | error +} + +const llmUsageSchema = ` +CREATE TABLE IF NOT EXISTS llm_usage ( + id BIGSERIAL PRIMARY KEY, + ts TIMESTAMPTZ NOT NULL DEFAULT now(), + task_id TEXT, + exploration_id BIGINT, + worker TEXT, + model TEXT, + profile_name TEXT, + latency_ms INTEGER, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read INTEGER NOT NULL DEFAULT 0, + cache_write INTEGER NOT NULL DEFAULT 0, + status TEXT +); +CREATE INDEX IF NOT EXISTS idx_llm_usage_task ON llm_usage(task_id); +CREATE INDEX IF NOT EXISTS idx_llm_usage_model ON llm_usage(task_id, model); +CREATE INDEX IF NOT EXISTS idx_llm_usage_exp ON llm_usage(exploration_id); +` + +// EnsureLLMUsageTable creates the llm_usage metering table if it does not exist. +func (d *DB) EnsureLLMUsageTable() error { + _, err := d.Exec(llmUsageSchema) + return err +} + +// InsertLLMUsage appends one metering row. Best-effort: callers log and continue on +// error (a lost metering row must never break the LLM call). +func (d *DB) InsertLLMUsage(u *LLMUsage) error { + var expID any + if u.ExplorationID > 0 { + expID = u.ExplorationID + } + _, err := d.Exec(` +INSERT INTO llm_usage(task_id, exploration_id, worker, model, profile_name, latency_ms, input_tokens, output_tokens, cache_read, cache_write, status) +VALUES (NULLIF($1,''),$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),$6,$7,$8,$9,$10,$11)`, + u.TaskID, expID, u.Worker, u.Model, u.ProfileName, + u.LatencyMs, u.InputTokens, u.OutputTokens, u.CacheRead, u.CacheWrite, u.Status) + return err +} + +// TokenByModel aggregates a task's LLM token usage grouped by model, most-used +// first, from the always-on llm_usage ledger. taskID is the task registry id. +// Accurate even with per-agent model bindings, pool rotation/failover, and +// interrupted runs, since every call (success or error) is metered. +func (d *DB) TokenByModel(taskID string) ([]ModelTokenStat, error) { + rows, err := d.Query(` +SELECT COALESCE(NULLIF(model,''),'(unknown)') AS model, COUNT(*) AS calls, + COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0), + COALESCE(SUM(cache_read),0), COALESCE(SUM(cache_write),0) +FROM llm_usage +WHERE COALESCE(task_id,'') = $1 +GROUP BY model +ORDER BY SUM(input_tokens) + SUM(output_tokens) DESC, model`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []ModelTokenStat{} + for rows.Next() { + var m ModelTokenStat + if err := rows.Scan(&m.Model, &m.Calls, &m.InputTokens, &m.OutputTokens, + &m.CacheReadTokens, &m.CacheWriteTokens); err != nil { + return nil, err + } + out = append(out, m) + } + return out, rows.Err() +} + +// ProfileUsage aggregates the whole ledger's token spend for one LLM profile +// (global, all tasks). Powers the dashboard's per-profile token card (new source). +type ProfileUsage struct { + ProfileName string `json:"profile_name"` + Calls int `json:"calls"` + Tasks int `json:"tasks"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheReadTokens int `json:"cache_read_tokens"` + CacheWriteTokens int `json:"cache_write_tokens"` +} + +// UsageByProfile returns global token spend grouped by profile name, most-used +// first. profile_name may be empty for calls made on env/non-persisted configs. +func (d *DB) UsageByProfile() ([]ProfileUsage, error) { + rows, err := d.Query(` +SELECT COALESCE(profile_name,'') AS profile_name, COUNT(*) AS calls, + COUNT(DISTINCT task_id) AS tasks, + COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0), + COALESCE(SUM(cache_read),0), COALESCE(SUM(cache_write),0) +FROM llm_usage +GROUP BY profile_name +ORDER BY SUM(input_tokens) + SUM(output_tokens) DESC`) + if err != nil { + return nil, err + } + defer rows.Close() + out := []ProfileUsage{} + for rows.Next() { + var p ProfileUsage + if err := rows.Scan(&p.ProfileName, &p.Calls, &p.Tasks, + &p.InputTokens, &p.OutputTokens, &p.CacheReadTokens, &p.CacheWriteTokens); err != nil { + return nil, err + } + out = append(out, p) + } + return out, rows.Err() +} + +// ProfileDayUsage is one (profile, UTC calendar day) token bucket for the daily +// chart. Unlike the activity-based chart, ts is the real call time, so this is +// actual per-day consumption rather than tokens bucketed by task creation date. +type ProfileDayUsage struct { + ProfileName string `json:"profile_name"` + Date string `json:"date"` // YYYY-MM-DD (UTC) + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheReadTokens int `json:"cache_read_tokens"` +} + +// UsageDaily returns per-(profile, day) token buckets for the past `days` days +// (default 365 when days<=0), so the dashboard can slice by profile + range. +func (d *DB) UsageDaily(days int) ([]ProfileDayUsage, error) { + if days <= 0 { + days = 365 + } + rows, err := d.Query(` +SELECT COALESCE(profile_name,'') AS profile_name, + to_char(ts AT TIME ZONE 'UTC', 'YYYY-MM-DD') AS day, + COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0), COALESCE(SUM(cache_read),0) +FROM llm_usage +WHERE ts >= now() - ($1 * interval '1 day') +GROUP BY profile_name, day +ORDER BY day`, days) + if err != nil { + return nil, err + } + defer rows.Close() + out := []ProfileDayUsage{} + for rows.Next() { + var p ProfileDayUsage + if err := rows.Scan(&p.ProfileName, &p.Date, &p.InputTokens, &p.OutputTokens, &p.CacheReadTokens); err != nil { + return nil, err + } + out = append(out, p) + } + return out, rows.Err() +} + +// ParseExpID turns the exploration-id segment parsed from a session string into an +// int64 (0 when empty/non-numeric, e.g. chat sessions keyed by conversation id). +func ParseExpID(s string) int64 { + n, err := strconv.ParseInt(s, 10, 64) + if err != nil || n < 0 { + return 0 + } + return n +} diff --git a/db/schema.sql b/db/schema.sql index 90779f7..fd55259 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -235,7 +235,10 @@ CREATE TABLE IF NOT EXISTS llm_profiles ( rate_per_second DOUBLE PRECISION NOT NULL DEFAULT 0, rate_per_minute DOUBLE PRECISION NOT NULL DEFAULT 0, context_window_k INTEGER NOT NULL DEFAULT 0, + -- 思考参数拆成两个独立字段：thinking_type=思考开关(''/disabled/enabled)， + -- reasoning_effort=思考强度(''/low/medium/high/xhigh/max)，互不牵连。 reasoning_effort TEXT NOT NULL DEFAULT '', + thinking_type TEXT NOT NULL DEFAULT '', is_default BOOLEAN NOT NULL DEFAULT false, -- 轮询(故障转移)参数，见 docs/LLM轮询设计.md： -- priority 顺位，越大越先被选中；激活配置(is_default)永远排链首，与本值无关。 @@ -253,6 +256,26 @@ CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS priority INTEGER NOT NULL DEFAULT 0; ALTER TABLE llm_profiles ADD COLUMN IF NOT EXISTS pool_exclude BOOLEAN NOT NULL DEFAULT false; +-- 思考开关字段 thinking_type，从旧的单一 reasoning_effort 语义一次性拆分而来。 +-- schema.sql 每次启动都执行，故迁移必须只跑一次：仅当该列尚不存在时才回填， +-- 否则每次启动都会把用户后来手动设的组合覆盖回去。旧 reasoning_effort 语义： +-- 'off' → 显式关闭 → thinking_type='disabled'，强度清空 +-- 'low/medium/high/max' → 开启+强度 → thinking_type='enabled'，强度保留 +-- '' → 不发送 → 两者皆空(默认) +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'llm_profiles' AND column_name = 'thinking_type' + ) THEN + ALTER TABLE llm_profiles ADD COLUMN thinking_type TEXT NOT NULL DEFAULT ''; + UPDATE llm_profiles SET thinking_type = 'enabled' + WHERE reasoning_effort IN ('low','medium','high','max'); + UPDATE llm_profiles SET thinking_type = 'disabled', reasoning_effort = '' + WHERE reasoning_effort = 'off'; + END IF; +END $$; + -- LLM 轮询熔断状态：某个配置连续失败(余额不足/key 失效/限流)后进入冷却，冷却期内 -- 轮询直接跳过它。内存态为准，这里落库只为重启后不丢冷却窗口——加载时只取尚未 -- 到期的行(open_until > now)，已到期的自然回到"正常"，等下一次调用半开试探。 @@ -391,7 +414,7 @@ CREATE TABLE IF NOT EXISTS agents ( llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL, current_prompt_id BIGINT, max_turns INTEGER NOT NULL DEFAULT 0, - run_seconds INTEGER NOT NULL DEFAULT 600, + run_seconds INTEGER NOT NULL DEFAULT 1200, web_search BOOLEAN NOT NULL DEFAULT false, interactive_shell BOOLEAN NOT NULL DEFAULT false, wrapup_prompt TEXT NOT NULL DEFAULT '', @@ -411,6 +434,8 @@ ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_merge_mode TEXT NOT NUL ALTER TABLE agents ADD COLUMN IF NOT EXISTS trigger_max_parallel INTEGER NOT NULL DEFAULT 5; -- per-agent LLM 绑定(agent 级默认模型):列自初版即在上方 CREATE 中,此 ALTER 仅为极旧库兜底(幂等)。 ALTER TABLE agents ADD COLUMN IF NOT EXISTS llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL; +-- run_seconds 单次 run 墙钟默认 600→1200:只改列默认(影响将来新插入的行),不动旧库存量行。 +ALTER TABLE agents ALTER COLUMN run_seconds SET DEFAULT 1200; CREATE INDEX IF NOT EXISTS idx_agents_llm_profile ON agents(llm_profile_id) WHERE llm_profile_id IS NOT NULL; DROP TRIGGER IF EXISTS trg_agents_upd ON agents; CREATE TRIGGER trg_agents_upd BEFORE UPDATE ON agents @@ -501,6 +526,26 @@ CREATE TABLE IF NOT EXISTS agent_skill_visibility ( ); CREATE INDEX IF NOT EXISTS idx_askv_skill ON agent_skill_visibility(skill_name); +-- Skill 调用账本（见 db/skill_usage.go）。一次 Skill() 调用一行，只记维度不记正文。 +-- 刻意不设外键：任务/会话删除后统计仍要保留（与 llm_usage 同理），skill 本身也只是 +-- 文件系统上的目录名，没有对应的表。 +CREATE TABLE IF NOT EXISTS skill_usage ( + id BIGSERIAL PRIMARY KEY, + ts TIMESTAMPTZ NOT NULL DEFAULT now(), + skill TEXT NOT NULL, + agent_key TEXT, + task_id BIGINT, + exploration_id BIGINT, + intent_id BIGINT, + session_id TEXT, + args_len INTEGER NOT NULL DEFAULT 0, + -- false = 模型点名了一个不存在的 skill(未命中)。这类行同样保留：它反映"想用但没有" + -- 的缺口，是补 skill 的依据。 + found BOOLEAN NOT NULL DEFAULT true +); +CREATE INDEX IF NOT EXISTS idx_skill_usage_skill ON skill_usage(skill, ts DESC); +CREATE INDEX IF NOT EXISTS idx_skill_usage_task ON skill_usage(task_id); + -- ===================================================================== -- H. 内置工具目录 -- ===================================================================== diff --git a/db/skill_usage.go b/db/skill_usage.go new file mode 100644 index 0000000..61d3a63 --- /dev/null +++ b/db/skill_usage.go @@ -0,0 +1,200 @@ +package db + +import ( + "database/sql" + "strings" + "time" +) + +// SkillUsage is one Skill() invocation — the always-on skill call ledger. Written +// from the Skill meta-tool's OnInvoke hook (server/assembly.go), one row per load. +// Carries only dimensions (which skill, which agent, which task/session), never the +// caller's args text — args_len is kept so an "empty vs substantial context" split +// is still possible without storing prompt content, mirroring llm_usage. +// +// Rows deliberately outlive their task: skill_usage has no foreign keys, so deleting +// a task keeps its skill statistics intact (same rationale as llm_usage). +type SkillUsage struct { + Skill string `json:"skill"` // skill directory name (matches agent_skill_visibility.skill_name) + AgentKey string `json:"agent_key"` // worker / planner / mainagent / custom agent key + TaskID int64 `json:"task_id"` // 0 for non-task runs (chat sessions) + ExplorationID int64 `json:"exploration_id"` // 0 when unknown + IntentID int64 `json:"intent_id"` // worker's intent node; 0 for planner/mainagent/chat + SessionID string `json:"session_id"` // chat conversation id; empty for task runs + ArgsLen int `json:"args_len"` + Found bool `json:"found"` // false = the model named a skill that does not exist +} + +// InsertSkillUsage appends one ledger row. Best-effort: callers log and continue on +// error (a lost metering row must never break a skill invocation). +func (d *DB) InsertSkillUsage(u *SkillUsage) error { + _, err := d.Exec(` +INSERT INTO skill_usage(skill, agent_key, task_id, exploration_id, intent_id, session_id, args_len, found) +VALUES ($1, NULLIF($2,''), $3, $4, $5, NULLIF($6,''), $7, $8)`, + u.Skill, u.AgentKey, nullIfZero(u.TaskID), nullIfZero(u.ExplorationID), + nullIfZero(u.IntentID), u.SessionID, u.ArgsLen, u.Found) + return err +} + +func nullIfZero(v int64) any { + if v > 0 { + return v + } + return nil +} + +// SkillStat is one skill's aggregate usage, for the skills page. +type SkillStat struct { + Skill string `json:"skill"` + Calls int `json:"calls"` + Tasks int `json:"tasks"` // distinct tasks that loaded it (chat runs excluded) + Agents []string `json:"agents"` // agent keys that loaded it, most-used first + LastUsed *time.Time `json:"last_used"` // nil when never called +} + +// SkillStats aggregates the whole ledger grouped by skill, most-used first. Skills +// that were never invoked are absent — callers merge against the skill list on disk. +// Only resolved calls count; misses are reported separately by MissingSkillStats. +func (d *DB) SkillStats() ([]SkillStat, error) { + // agent keys come back as one comma-joined string rather than text[]: the pgx + // stdlib driver has no database/sql Scan target for arrays, and agent keys are + // [a-z0-9_-] so a comma join is unambiguous. + rows, err := d.Query(` +SELECT skill, COUNT(*) AS calls, + COUNT(DISTINCT task_id) AS tasks, + COALESCE(STRING_AGG(DISTINCT agent_key, ','), '') AS agents, + MAX(ts) AS last_used +FROM skill_usage +WHERE found +GROUP BY skill +ORDER BY COUNT(*) DESC, skill`) + if err != nil { + return nil, err + } + defer rows.Close() + out := []SkillStat{} + for rows.Next() { + var s SkillStat + var agents string + var lastUsed sql.NullTime + if err := rows.Scan(&s.Skill, &s.Calls, &s.Tasks, &agents, &lastUsed); err != nil { + return nil, err + } + s.Agents = []string{} + if agents != "" { + s.Agents = strings.Split(agents, ",") + } + if lastUsed.Valid { + t := lastUsed.Time + s.LastUsed = &t + } + out = append(out, s) + } + return out, rows.Err() +} + +// MissingSkillStats returns the skill names agents asked for that do not exist, +// most-requested first — the "wished it existed" gap list. Names come from the model +// so they are shown as-is (already length-capped at insert time). +func (d *DB) MissingSkillStats(limit int) ([]SkillStat, error) { + if limit <= 0 || limit > 200 { + limit = 20 + } + rows, err := d.Query(` +SELECT skill, COUNT(*) AS calls, + COALESCE(STRING_AGG(DISTINCT agent_key, ','), '') AS agents, + MAX(ts) AS last_used +FROM skill_usage +WHERE NOT found +GROUP BY skill +ORDER BY COUNT(*) DESC, skill +LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer rows.Close() + out := []SkillStat{} + for rows.Next() { + var s SkillStat + var agents string + var lastUsed sql.NullTime + if err := rows.Scan(&s.Skill, &s.Calls, &agents, &lastUsed); err != nil { + return nil, err + } + s.Agents = []string{} + if agents != "" { + s.Agents = strings.Split(agents, ",") + } + if lastUsed.Valid { + t := lastUsed.Time + s.LastUsed = &t + } + out = append(out, s) + } + return out, rows.Err() +} + +// SkillCall is one row of a skill's recent-call list (detail panel). +type SkillCall struct { + TS time.Time `json:"ts"` + AgentKey string `json:"agent_key"` + TaskID int64 `json:"task_id"` + SessionID string `json:"session_id"` + ArgsLen int `json:"args_len"` +} + +// RecentSkillCalls returns the most recent invocations of one skill, newest first. +func (d *DB) RecentSkillCalls(skill string, limit int) ([]SkillCall, error) { + if limit <= 0 || limit > 200 { + limit = 50 + } + rows, err := d.Query(` +SELECT ts, COALESCE(agent_key,''), COALESCE(task_id,0), COALESCE(session_id,''), args_len +FROM skill_usage +WHERE skill = $1 AND found +ORDER BY ts DESC +LIMIT $2`, skill, limit) + if err != nil { + return nil, err + } + defer rows.Close() + out := []SkillCall{} + for rows.Next() { + var c SkillCall + if err := rows.Scan(&c.TS, &c.AgentKey, &c.TaskID, &c.SessionID, &c.ArgsLen); err != nil { + return nil, err + } + out = append(out, c) + } + return out, rows.Err() +} + +// SkillCallsByTask counts a task's skill loads, most-used first. Powers a per-task +// view of which procedures its agents actually reached for. +func (d *DB) SkillCallsByTask(taskID int64) ([]SkillStat, error) { + rows, err := d.Query(` +SELECT skill, COUNT(*) AS calls, MAX(ts) AS last_used +FROM skill_usage +WHERE task_id = $1 AND found +GROUP BY skill +ORDER BY COUNT(*) DESC, skill`, taskID) + if err != nil { + return nil, err + } + defer rows.Close() + out := []SkillStat{} + for rows.Next() { + var s SkillStat + var lastUsed sql.NullTime + if err := rows.Scan(&s.Skill, &s.Calls, &lastUsed); err != nil { + return nil, err + } + s.Agents = []string{} + if lastUsed.Valid { + t := lastUsed.Time + s.LastUsed = &t + } + out = append(out, s) + } + return out, rows.Err() +} diff --git a/db/skill_usage_test.go b/db/skill_usage_test.go new file mode 100644 index 0000000..4625696 --- /dev/null +++ b/db/skill_usage_test.go @@ -0,0 +1,116 @@ +package db + +import ( + "testing" +) + +// TestSkillUsageLedger writes a few rows and checks the three aggregates the +// skills page relies on: per-skill totals (resolved calls only), the miss list, +// and the recent-call list. Rows are namespaced by a unique skill name so the +// shared dev DB stays usable and the test cleans up after itself. +func TestSkillUsageLedger(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + const ( + skillA = "zz-test-skill-a" + missing = "zz-test-skill-missing" + ) + cleanup := func() { + _, _ = d.Exec(`DELETE FROM skill_usage WHERE skill IN ($1,$2)`, skillA, missing) + } + cleanup() + defer cleanup() + + rows := []*SkillUsage{ + {Skill: skillA, AgentKey: "worker", TaskID: 991, ExplorationID: 5, IntentID: 7, ArgsLen: 12, Found: true}, + {Skill: skillA, AgentKey: "worker", TaskID: 991, ExplorationID: 5, ArgsLen: 0, Found: true}, + {Skill: skillA, AgentKey: "planner", TaskID: 992, ExplorationID: 6, Found: true}, + {Skill: skillA, AgentKey: "chatbot", SessionID: "conv-1", Found: true}, + {Skill: missing, AgentKey: "worker", TaskID: 991, Found: false}, + {Skill: missing, AgentKey: "worker", TaskID: 991, Found: false}, + } + for _, r := range rows { + if err := d.InsertSkillUsage(r); err != nil { + t.Fatalf("insert %s: %v", r.Skill, err) + } + } + + stats, err := d.SkillStats() + if err != nil { + t.Fatal(err) + } + var got *SkillStat + for i := range stats { + if stats[i].Skill == skillA { + got = &stats[i] + } + if stats[i].Skill == missing { + t.Fatalf("misses must not appear in SkillStats: %+v", stats[i]) + } + } + if got == nil { + t.Fatalf("skill %s absent from SkillStats", skillA) + } + if got.Calls != 4 { + t.Errorf("calls: want 4, got %d", got.Calls) + } + // 991 + 992; the chat row has a NULL task_id and COUNT(DISTINCT) skips NULLs. + if got.Tasks != 2 { + t.Errorf("tasks: want 2, got %d", got.Tasks) + } + if len(got.Agents) != 3 { + t.Errorf("agents: want 3 distinct, got %v", got.Agents) + } + if got.LastUsed == nil { + t.Error("last_used must be set") + } + + miss, err := d.MissingSkillStats(10) + if err != nil { + t.Fatal(err) + } + var missCalls int + for _, m := range miss { + if m.Skill == missing { + missCalls = m.Calls + } + } + if missCalls != 2 { + t.Errorf("missing calls: want 2, got %d", missCalls) + } + + calls, err := d.RecentSkillCalls(skillA, 10) + if err != nil { + t.Fatal(err) + } + if len(calls) != 4 { + t.Fatalf("recent calls: want 4, got %d", len(calls)) + } + // newest first + for i := 1; i < len(calls); i++ { + if calls[i].TS.After(calls[i-1].TS) { + t.Errorf("recent calls not newest-first at %d", i) + } + } + + byTask, err := d.SkillCallsByTask(991) + if err != nil { + t.Fatal(err) + } + var taskCalls int + for _, s := range byTask { + if s.Skill == skillA { + taskCalls = s.Calls + } + if s.Skill == missing { + t.Errorf("misses must not appear in SkillCallsByTask: %+v", s) + } + } + if taskCalls != 2 { + t.Errorf("task 991 calls: want 2, got %d", taskCalls) + } +} diff --git a/db/task_scope.go b/db/task_scope.go index 2c69dcd..af1fc4a 100644 --- a/db/task_scope.go +++ b/db/task_scope.go @@ -113,11 +113,15 @@ func (s *AssetStore) AddAutoScope(taskID int64, assetType, domain, rawURL, ip st return nil } -// AddAgentScope parses an agent-provided (kind,value) and records it (source='agent'). +// AddAgentScope parses a (kind,value) pair and records it as task scope. +// source is "agent" when called from an LLM tool, "manual" from the UI. // company: value = company name or id (must already exist). root_domain/subdomain: // value = a domain. ip/cidr: value = an IP or CIDR (bare IP → /32,/128). -func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (TaskScope, error) { - ts := TaskScope{TaskID: taskID, Kind: kind, Source: "agent", Reason: reason} +func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason, source string) (TaskScope, error) { + if source == "" { + source = "agent" + } + ts := TaskScope{TaskID: taskID, Kind: kind, Source: source, Reason: reason} if taskID <= 0 { return ts, fmt.Errorf("需要 task_id") } @@ -184,6 +188,17 @@ func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (Ta return ts, nil } +// DeleteTaskScope removes a single scope row by id, scoped to the given task. +// Returns whether a row was actually deleted. +func (s *AssetStore) DeleteTaskScope(taskID, scopeID int64) (bool, error) { + res, err := s.db.Exec(`DELETE FROM task_scope WHERE id=$1 AND task_id=$2`, scopeID, taskID) + if err != nil { + return false, err + } + n, _ := res.RowsAffected() + return n > 0, nil +} + // ListTaskScope returns all scope rows for a task. func (s *AssetStore) ListTaskScope(taskID int64) ([]TaskScope, error) { rows, err := s.db.Query(` diff --git a/db/tasks.go b/db/tasks.go index e5c4dc9..3740ccc 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -437,6 +437,8 @@ DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID) } result.LLMRecordsDeleted, _ = res.RowsAffected() } + // llm_usage (the token metering ledger) is intentionally NOT deleted with the + // task — it is kept as historical accounting even after the task is gone. if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil { return result, err } diff --git a/llmrec/llmrec.go b/llmrec/llmrec.go index cf0be46..1bbd55c 100644 --- a/llmrec/llmrec.go +++ b/llmrec/llmrec.go @@ -11,9 +11,9 @@ import ( "strings" "time" + "github.com/Autumn-27/artex/db" "github.com/Autumn-27/norma/llm" "github.com/Autumn-27/norma/transcript" - "github.com/Autumn-27/artex/db" ) type taskIDContextKey struct{} @@ -41,16 +41,25 @@ func TaskIDFrom(ctx context.Context) string { type Recorder struct { inner llm.Provider pg *db.DB - model string // model name (from config, not in CompletionRequest) - prof string // LLM profile name (from llm_profiles) - enabled func() bool // reports whether recording is currently on; nil = always record + model string // model name (from config, not in CompletionRequest) + prof string // LLM profile name (from llm_profiles) + // thinkingType / reasoningEffort 是配置级思考参数(思考开关 / 思考强度)。它们在 + // norma 的 buildBody() 里从 provider 配置注入真正的 HTTP body,不出现在 + // CompletionRequest 上,故 Recorder 需在此单独带一份,序列化时写进录制。 + thinkingType string + reasoningEffort string + enabled func() bool // reports whether recording is currently on; nil = always record } // Wrap returns a Provider that records calls to pg when enabled() reports true. -// profName is the LLM profile name (e.g. "default"); may be empty. A nil enabled -// predicate records unconditionally (backward-compatible default). -func Wrap(inner llm.Provider, pg *db.DB, model, profName string, enabled func() bool) *Recorder { - return &Recorder{inner: inner, pg: pg, model: model, prof: profName, enabled: enabled} +// profName is the LLM profile name (e.g. "default"); may be empty. thinkingType / +// reasoningEffort are the config-level thinking params actually sent to the API +// (empty = not sent). A nil enabled predicate records unconditionally. +func Wrap(inner llm.Provider, pg *db.DB, model, profName, thinkingType, reasoningEffort string, enabled func() bool) *Recorder { + return &Recorder{ + inner: inner, pg: pg, model: model, prof: profName, + thinkingType: thinkingType, reasoningEffort: reasoningEffort, enabled: enabled, + } } // parseSession extracts task id and worker role from session strings like @@ -84,14 +93,15 @@ func parseSession(s string) (taskID, worker string) { // harness via transcript.WithSessionID; reading it per-call is race-free even when // planner and multiple workers share one Recorder instance. func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { - // Recording off → passthrough with zero overhead: no request serialization, - // no response accumulation, no DB insert. - if r.enabled != nil && !r.enabled() { - return r.inner.Stream(ctx, req) - } + // Body recording (the heavy debug trace: full request/response) is gated by the + // llm_record setting. Lightweight usage metering always runs — it powers token + // stats and must be complete even for interrupted/failed runs, so it is NOT + // gated. Only body serialization + response accumulation are skipped when off. + recordBodies := r.enabled == nil || r.enabled() start := time.Now() session := transcript.SessionIDFrom(ctx) parsedID, worker := parseSession(session) + expID := db.ParseExpID(parsedID) taskID := TaskIDFrom(ctx) if taskID == "" { // Backward compatibility for non-task callers. For task calls, the task @@ -99,51 +109,101 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S taskID = parsedID } - // Serialize the request for storage. - reqBody := serializeRequest(req) + // Serialize the request only when storing bodies (this is the expensive part). + reqBody := "" + if recordBodies { + reqBody = r.serializeRequest(req) + } return func(yield func(llm.StreamEvent, error) bool) { var ( - textBuf strings.Builder + textBuf strings.Builder thinkingBuf strings.Builder - usage llm.Usage - stopReason string - streamErr error + usage llm.Usage + stopReason string + streamErr error ) + finish := func(err error) { + status := "ok" + if err != nil { + status = "error" + } + // Lightweight metering row — always written. + r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) + // Heavy trace row — only when body recording is on. + if recordBodies { + r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err) + } + } + for ev, err := range r.inner.Stream(ctx, req) { if err != nil { streamErr = err - // Record the failed call, then propagate. - r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, "", err) + finish(streamErr) if !yield(ev, err) { return } return } - // Accumulate response content. + // Always track usage (cheap); accumulate text/thinking only for bodies. + // Anthropic (and the other providers) split token usage across events: + // message_start carries ONLY the input side (input + cache), message_delta + // ONLY the output. They must be FOLDED with Add — overwriting on delta + // would zero out the input already counted at start (mirrors the SDK's own + // llm.Accumulator; see norma/llm/accumulate.go). switch ev.Type { case llm.SETextDelta: - textBuf.WriteString(ev.Text) + if recordBodies { + textBuf.WriteString(ev.Text) + } case llm.SEThinkingDelta: - thinkingBuf.WriteString(ev.Text) + if recordBodies { + thinkingBuf.WriteString(ev.Text) + } + case llm.SEMessageStart: + usage.Add(ev.Usage) case llm.SEMessageDelta: if ev.StopReason != "" { stopReason = ev.StopReason } - usage = ev.Usage - case llm.SEMessageStart: - if ev.Usage.InputTokens > 0 { - usage = ev.Usage - } + usage.Add(ev.Usage) } if !yield(ev, err) { return } } - // Stream completed normally — record it. - r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, streamErr) + // Stream completed normally. + finish(streamErr) + } +} + +// recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips +// zero-token calls with no model, which carry nothing worth metering. +func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) { + if r.pg == nil { + return + } + if r.model == "" && usage.InputTokens == 0 && usage.OutputTokens == 0 && + usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 { + return + } + err := r.pg.InsertLLMUsage(&db.LLMUsage{ + TaskID: taskID, + ExplorationID: expID, + Worker: worker, + Model: r.model, + ProfileName: r.prof, + LatencyMs: latencyMs, + InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, + CacheRead: usage.CacheReadTokens, + CacheWrite: usage.CacheWriteTokens, + Status: status, + }) + if err != nil { + log.Printf("[llmusage] insert: %v", err) } } @@ -204,12 +264,22 @@ func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, re } // serializeRequest builds a JSON representation of the completion request. -func serializeRequest(req llm.CompletionRequest) string { +func (r *Recorder) serializeRequest(req llm.CompletionRequest) string { m := map[string]any{ "system": req.System, "messages": req.Messages, "max_tokens": req.MaxTokens, } + // 记录本次调用实际发出的思考参数。type 采用「有效值」：每请求覆盖 req.Thinking + // 优先于配置级 thinkingType(与 norma buildBody 的判定一致，如 compaction 摘要会 + // 强制 disabled)；effort 无每请求覆盖，直接取配置值。两者皆空则不写 thinking 字段。 + effType := r.thinkingType + if req.Thinking != "" { + effType = req.Thinking + } + if effType != "" || r.reasoningEffort != "" { + m["thinking"] = map[string]string{"type": effType, "effort": r.reasoningEffort} + } if len(req.Tools) > 0 { // Store tool names only (full schemas are huge). names := make([]string, len(req.Tools)) diff --git a/report/findings.go b/report/findings.go new file mode 100644 index 0000000..265bfd3 --- /dev/null +++ b/report/findings.go @@ -0,0 +1,165 @@ +package report + +import ( + "bytes" + "encoding/csv" + "fmt" + "path" + "regexp" + "sort" + "strings" + "time" + + "github.com/Autumn-27/artex/db" +) + +// 发现页「导出」用的渲染:把一批 findings 表行渲染成汇总 Markdown、单条 Markdown、 +// 或 CSV。JSON 由 server 层直接用 DTO 序列化,不在此处。 + +// sortFindingsForExport 按严重等级降序、再按时间倒序排,与汇总报告的分组一致。 +func sortFindingsForExport(fs []*db.DBFinding) { + sort.SliceStable(fs, func(i, j int) bool { + ri, rj := sevRank[fs[i].Severity], sevRank[fs[j].Severity] + if ri != rj { + return ri < rj // sevRank 越小越严重 + } + return fs[i].CreatedAt.After(fs[j].CreatedAt) + }) +} + +// findingTitle 取漏洞可读标题:名称 → 类别 → 「未分类」。 +func findingTitle(f *db.DBFinding) string { + return nz(f.Name, nz(f.VulnClass, "未分类")) +} + +// FindingsMarkdown 把一批 findings 整合成一份汇总报告(摘要 + 按严重等级分组, +// 每条含类别/状态/所属任务/证据/详细报告)。 +func FindingsMarkdown(fs []*db.DBFinding, generatedAt time.Time) string { + items := append([]*db.DBFinding(nil), fs...) + sortFindingsForExport(items) + + var b strings.Builder + b.WriteString("# 漏洞发现汇总报告\n\n") + fmt.Fprintf(&b, "- **生成时间**：%s\n", generatedAt.Format("2006-01-02 15:04:05")) + fmt.Fprintf(&b, "- **发现总数**：%d 个\n\n", len(items)) + + // 摘要:各严重等级计数。 + counts := map[string]int{} + for _, f := range items { + counts[f.Severity]++ + } + b.WriteString("## 摘要\n\n") + b.WriteString("| 严重等级 | 数量 |\n| --- | --- |\n") + for _, s := range []struct{ key, label string }{ + {"critical", "严重"}, {"high", "高危"}, {"medium", "中危"}, {"low", "低危"}, + } { + fmt.Fprintf(&b, "| %s | %d |\n", s.label, counts[s.key]) + } + b.WriteString("\n") + + if len(items) == 0 { + b.WriteString("_无匹配的漏洞。_\n") + return b.String() + } + + b.WriteString("## 漏洞明细\n\n") + for i, f := range items { + fmt.Fprintf(&b, "### %d. [%s] %s\n\n", i+1, strings.ToUpper(nz(f.Severity, "info")), findingTitle(f)) + if f.VulnClass != "" { + fmt.Fprintf(&b, "- **类别**：%s\n", f.VulnClass) + } + fmt.Fprintf(&b, "- **状态**：%s\n", nz(f.Status, "pending")) + if desc := strings.TrimSpace(f.TaskDescription); desc != "" { + fmt.Fprintf(&b, "- **所属任务**：%s\n", desc) + } + fmt.Fprintf(&b, "- **发现时间**：%s\n\n", f.CreatedAt.Format("2006-01-02 15:04:05")) + if s := strings.TrimSpace(f.Summary); s != "" { + fmt.Fprintf(&b, "%s\n\n", s) + } + if e := strings.TrimSpace(f.Evidence); e != "" { + fmt.Fprintf(&b, "**证据：**\n\n```\n%s\n```\n\n", e) + } + if rep := strings.TrimSpace(f.Report); rep != "" { + b.WriteString("**详细报告：**\n\n") + b.WriteString(rep) + b.WriteString("\n\n") + } + b.WriteString("---\n\n") + } + return b.String() +} + +// SingleFindingMarkdown 渲染单条漏洞为一份独立 Markdown(用于「一漏洞一文件」打包)。 +func SingleFindingMarkdown(f *db.DBFinding, generatedAt time.Time) string { + var b strings.Builder + fmt.Fprintf(&b, "# [%s] %s\n\n", strings.ToUpper(nz(f.Severity, "info")), findingTitle(f)) + if f.VulnClass != "" { + fmt.Fprintf(&b, "- **类别**：%s\n", f.VulnClass) + } + fmt.Fprintf(&b, "- **严重等级**：%s\n", nz(f.Severity, "info")) + fmt.Fprintf(&b, "- **状态**：%s\n", nz(f.Status, "pending")) + if desc := strings.TrimSpace(f.TaskDescription); desc != "" { + fmt.Fprintf(&b, "- **所属任务**：%s\n", desc) + } + fmt.Fprintf(&b, "- **发现时间**：%s\n", f.CreatedAt.Format("2006-01-02 15:04:05")) + fmt.Fprintf(&b, "- **生成时间**：%s\n\n", generatedAt.Format("2006-01-02 15:04:05")) + if s := strings.TrimSpace(f.Summary); s != "" { + fmt.Fprintf(&b, "## 概述\n\n%s\n\n", s) + } + if e := strings.TrimSpace(f.Evidence); e != "" { + fmt.Fprintf(&b, "## 证据\n\n```\n%s\n```\n\n", e) + } + if rep := strings.TrimSpace(f.Report); rep != "" { + b.WriteString("## 详细报告\n\n") + b.WriteString(rep) + b.WriteString("\n") + } + return b.String() +} + +var unsafeFilenameChars = regexp.MustCompile(`[^\p{Han}\p{L}\p{N}._-]+`) + +// FindingFilename 为「一漏洞一文件」生成安全的 .md 文件名,形如 +// `critical_SQL注入_#123.md`。去掉路径分隔符与控制字符,避免 zip 内非法路径。 +func FindingFilename(f *db.DBFinding) string { + sev := nz(f.Severity, "info") + title := findingTitle(f) + name := fmt.Sprintf("%s_%s_#%d", sev, title, f.ID) + name = unsafeFilenameChars.ReplaceAllString(name, "_") + name = strings.Trim(name, "._") + if name == "" { + name = fmt.Sprintf("finding_%d", f.ID) + } + // 防御性:再剥一层路径,杜绝 zip slip。 + name = path.Base(name) + if len(name) > 120 { + name = name[:120] + } + return name + ".md" +} + +// FindingsCSV 把一批 findings 渲染成 CSV(带 UTF-8 BOM,便于 Excel 正确识别中文)。 +// 不含大段 report/evidence 全文,只放摘要类字段;需要全文用 Markdown/JSON 导出。 +func FindingsCSV(fs []*db.DBFinding) []byte { + items := append([]*db.DBFinding(nil), fs...) + sortFindingsForExport(items) + + var buf bytes.Buffer + buf.WriteString("\xEF\xBB\xBF") // UTF-8 BOM + w := csv.NewWriter(&buf) + _ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述"}) + for _, f := range items { + _ = w.Write([]string{ + fmt.Sprintf("%d", f.ID), + findingTitle(f), + f.VulnClass, + nz(f.Severity, "info"), + nz(f.Status, "pending"), + f.TaskDescription, + f.CreatedAt.Format("2006-01-02 15:04:05"), + strings.TrimSpace(f.Summary), + }) + } + w.Flush() + return buf.Bytes() +} diff --git a/server/assembly.go b/server/assembly.go index 49fc162..426a9b6 100644 --- a/server/assembly.go +++ b/server/assembly.go @@ -137,7 +137,11 @@ func wireAgentAugment(pg *db.DB, skillDir string, hostTools func() ([]actool.Cor } return actool.RenderDeferredToolsBlock(reveal) } - extra = append(extra, reg.Tool()) + // Attribution for the usage ledger: ToolAugment only gets (ctx, agentKey), + // so the run's task/session ids ride in on the ctx (agent.RunInfo). Read it + // once here — this closure is rebuilt per run, so the captured value always + // belongs to this run. + extra = append(extra, meterSkillTool(reg.Tool(), pg, reg, a.Key, agent.RunInfoFrom(ctx))) } // host tools (traffic / orchestration / custom) — added to every agent's base; diff --git a/server/commands.go b/server/commands.go index bcb35e0..12904a7 100644 --- a/server/commands.go +++ b/server/commands.go @@ -4,6 +4,8 @@ import ( "net/http" "strconv" "strings" + + "github.com/Autumn-27/artex/db" ) // pgListCommands returns tool executions (any tool) from the activity table. @@ -45,6 +47,53 @@ func (s *Server) pgListLLMRecords(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"records": records, "total": total}) } +// pgTokenByModel returns a task's LLM token usage grouped by model, from the +// always-on llm_usage metering ledger. Accurate even with per-agent model bindings, +// pool rotation/failover, and interrupted runs (every call is metered on success or +// error). No recording toggle required. +func (s *Server) pgTokenByModel(w http.ResponseWriter, r *http.Request) { + id := strings.TrimSpace(r.URL.Query().Get("task")) + if id == "" { + writeErr(w, 400, "missing task") + return + } + pg := s.m.PG() + if pg == nil { + writeJSON(w, 200, map[string]any{"models": []db.ModelTokenStat{}}) + return + } + models, err := pg.TokenByModel(id) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"models": models}) +} + +// pgUsageStats returns global llm_usage aggregates for the dashboard's "new" token +// view: per-profile totals (all-time) + per-(profile, day) buckets for the daily +// chart. Sourced from the always-on metering ledger, so it is accurate across +// per-agent bindings, pool rotation, and interrupted runs. +func (s *Server) pgUsageStats(w http.ResponseWriter, r *http.Request) { + pg := s.m.PG() + if pg == nil { + writeJSON(w, 200, map[string]any{"by_profile": []db.ProfileUsage{}, "daily": []db.ProfileDayUsage{}}) + return + } + days := atoiDefault(r.URL.Query().Get("days"), 365) + byProfile, err := pg.UsageByProfile() + if err != nil { + writeErr(w, 500, err.Error()) + return + } + daily, err := pg.UsageDaily(days) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"by_profile": byProfile, "daily": daily}) +} + // pgLLMTasks returns distinct recorded tasks with counts, for the page's task // picker (pick a task → filter the list, then delete its conversations). func (s *Server) pgLLMTasks(w http.ResponseWriter, r *http.Request) { diff --git a/server/dto.go b/server/dto.go index 9c320c0..599e47f 100644 --- a/server/dto.go +++ b/server/dto.go @@ -500,6 +500,7 @@ type LLMProfileDTO struct { RatePerSecond float64 `json:"rate_per_second"` RatePerMinute float64 `json:"rate_per_minute"` ContextWindowK int `json:"context_window_k"` + ThinkingType string `json:"thinking_type"` ReasoningEffort string `json:"reasoning_effort"` IsDefault bool `json:"is_default"` // 轮询(故障转移)参数：priority 越大越先被选中(激活配置恒为链首)； @@ -520,6 +521,7 @@ func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO { RatePerSecond: p.RatePerSecond, RatePerMinute: p.RatePerMinute, ContextWindowK: p.ContextWindowK, + ThinkingType: p.ThinkingType, ReasoningEffort: p.ReasoningEffort, IsDefault: p.IsDefault, Priority: p.Priority, diff --git a/server/engine.go b/server/engine.go index c52e76b..4606a8f 100644 --- a/server/engine.go +++ b/server/engine.go @@ -926,7 +926,8 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker return true } // if a pause cancelled this run mid-flight, return the intent to the frontier - // so it is re-claimed (and re-run from scratch) on resume — not marked done. + // 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") return true diff --git a/server/manager.go b/server/manager.go index 5f8dc36..7779088 100644 --- a/server/manager.go +++ b/server/manager.go @@ -250,6 +250,9 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { if err := pg.EnsureLLMRecordsTable(); err != nil { log.Printf("[llmrec] create table: %v", err) } + if err := pg.EnsureLLMUsageTable(); err != nil { + log.Printf("[llmusage] create table: %v", err) + } m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)} if proxyAddr != "" { tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr) diff --git a/server/server.go b/server/server.go index 2f44d26..0b82c82 100644 --- a/server/server.go +++ b/server/server.go @@ -1,6 +1,7 @@ package server import ( + "archive/zip" "context" "encoding/json" @@ -274,14 +275,14 @@ func (s *Server) agentMaxTurns(key string) int { } // agentRunSeconds returns the configured wall-clock run budget (seconds) for an -// agent key (0 = unlimited; 600 fallback when no DB or no row, matching schema). +// agent key (0 = unlimited; 1200 fallback when no DB or no row, matching schema). func (s *Server) agentRunSeconds(key string) int { if s.m.pg == nil { - return 600 + return 1200 } a, err := s.m.pg.GetAgentByKey(key) if err != nil || a == nil { - return 600 + return 1200 } return a.RunSecs } @@ -295,6 +296,7 @@ func (s *Server) loadLLMConfig() (agent.Config, bool) { cfg := agent.ConfigFrom(p.Format, p.Model, p.BaseURL, p.APIKey, p.Proxy) cfg.RatePerSecond, cfg.RatePerMinute = p.RatePerSecond, p.RatePerMinute cfg.ContextWindowK = p.ContextWindowK + cfg.ThinkingType = p.ThinkingType cfg.ReasoningEffort = p.ReasoningEffort if cfg.APIKey == "" { return cfg, false @@ -328,7 +330,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error { newID, err := s.m.pg.SaveProfile(&db.LLMProfile{ ID: id, Name: "default", Format: format, Model: cfg.Model, BaseURL: cfg.BaseURL, Proxy: cfg.Proxy, APIKey: cfg.APIKey, RatePerSecond: cfg.RatePerSecond, RatePerMinute: cfg.RatePerMinute, - ContextWindowK: cfg.ContextWindowK, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, + ContextWindowK: cfg.ContextWindowK, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, Priority: priority, PoolExclude: poolExclude, }) if err != nil { @@ -408,7 +410,7 @@ func (s *Server) applyLLM(cfg agent.Config) error { s.cfgMu.Lock() profName := s.llmProf s.cfgMu.Unlock() - prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, s.m.LLMRecordEnabled) + prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled) } // LLM 轮询(默认关):把激活配置包进故障转移链,当前配置不可用时自动切下一个。 // 只影响「走全局激活配置」的这条路径——agent 绑定 / 任务 pin 的走 providerForProfile, @@ -458,6 +460,7 @@ func (s *Server) loadProfileConfig(id int64) (agent.Config, bool) { cfg := agent.ConfigFrom(p.Format, p.Model, p.BaseURL, p.APIKey, p.Proxy) cfg.RatePerSecond, cfg.RatePerMinute = p.RatePerSecond, p.RatePerMinute cfg.ContextWindowK = p.ContextWindowK + cfg.ThinkingType = p.ThinkingType cfg.ReasoningEffort = p.ReasoningEffort if cfg.APIKey == "" { return cfg, false @@ -501,7 +504,7 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) } // Wrap with recorder, tagged with this profile's name. if p, _ := s.m.pg.ProfileByID(id); p != nil { - prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, s.m.LLMRecordEnabled) + prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled) } s.provCacheMu.Lock() if generation != s.provCacheGen { @@ -647,6 +650,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/workspace/download", s.wsDownload) mux.HandleFunc("POST /api/workspace/upload", s.wsUpload) 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("POST /api/tasks/{id}/control", s.control) mux.HandleFunc("PUT /api/tasks/{id}/llm", s.updateTaskLLMProfiles) mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/control", s.controlIntent) @@ -675,6 +680,7 @@ 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/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("PATCH /api/exploration/findings/{id}", s.patchFinding) @@ -688,6 +694,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/exploration/tokens", s.tokenStats) mux.HandleFunc("GET /api/tokens/daily", s.tokenDailyStats) mux.HandleFunc("GET /api/tokens/conversations", s.conversationTokens) + mux.HandleFunc("GET /api/tokens/usage", s.pgUsageStats) // 全局 llm_usage 聚合（仪表盘新版视图） mux.HandleFunc("GET /api/audit", s.getAudit) mux.HandleFunc("POST /api/gc", s.gc) @@ -700,6 +707,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/llm/records", s.pgListLLMRecords) mux.HandleFunc("DELETE /api/llm/records", s.pgDeleteLLMRecords) mux.HandleFunc("GET /api/llm/records/tasks", s.pgLLMTasks) + mux.HandleFunc("GET /api/llm/records/by-model", s.pgTokenByModel) // 按模型聚合本任务 token 用量 mux.HandleFunc("GET /api/llm/records/{id}", s.pgGetLLMRecord) mux.HandleFunc("GET /api/settings", s.getSettings) mux.HandleFunc("PUT /api/settings", s.putSettings) @@ -770,6 +778,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /api/skills", s.fsCreateSkill) mux.HandleFunc("POST /api/skills/upload", s.fsUploadSkill) mux.HandleFunc("DELETE /api/skills/{name}", s.fsDeleteSkill) + mux.HandleFunc("GET /api/skills/missing", s.fsMissingSkills) // 未命中(想调但不存在)的 skill 名 + mux.HandleFunc("GET /api/skills/{name}/usage", s.fsSkillUsage) // 单个 skill 的最近调用 mux.HandleFunc("PUT /api/skills/{name}/meta", s.fsUpdateSkillMeta) mux.HandleFunc("POST /api/skills/{name}/dirs", s.fsCreateDir) mux.HandleFunc("GET /api/skills/{name}/files", s.fsListFiles) @@ -994,7 +1004,9 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) { 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. New turns are rejected while the task is paused. + // 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 @@ -1167,6 +1179,7 @@ func (s *Server) getLLM(w http.ResponseWriter, r *http.Request) { "rate_per_second": s.llmCfg.RatePerSecond, "rate_per_minute": s.llmCfg.RatePerMinute, "context_window_k": s.llmCfg.ContextWindowK, + "thinking_type": s.llmCfg.ThinkingType, "reasoning_effort": s.llmCfg.ReasoningEffort, }) } @@ -1182,6 +1195,7 @@ func (s *Server) setLLM(w http.ResponseWriter, r *http.Request) { RatePerSecond float64 `json:"rate_per_second"` RatePerMinute float64 `json:"rate_per_minute"` ContextWindowK int `json:"context_window_k"` + ThinkingType string `json:"thinking_type"` ReasoningEffort string `json:"reasoning_effort"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { @@ -1190,6 +1204,7 @@ func (s *Server) setLLM(w http.ResponseWriter, r *http.Request) { } cfg := agent.ConfigFrom(req.Provider, req.Model, req.BaseURL, req.APIKey, req.Proxy) cfg.RatePerSecond, cfg.RatePerMinute = req.RatePerSecond, req.RatePerMinute + cfg.ThinkingType = req.ThinkingType cfg.ReasoningEffort = req.ReasoningEffort if k := req.ContextWindowK; k > 0 { // 0 = keep default (200K); cap at 1M if k > 1000 { @@ -1236,6 +1251,7 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { BaseURL string `json:"base_url"` Proxy string `json:"proxy"` APIKey string `json:"api_key"` + ThinkingType string `json:"thinking_type"` ReasoningEffort string `json:"reasoning_effort"` ProfileID *int64 `json:"profile_id"` // 测已存 profile 时传入：api_key 为空则用它存的 key } @@ -1246,6 +1262,7 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { cfg := agent.ConfigFrom(req.Provider, req.Model, req.BaseURL, req.APIKey, req.Proxy) // mirror production: send the SAME thinking params so a provider that rejects the // reasoning_effort/thinking field fails the test too (no false "test ok, run 400"). + cfg.ThinkingType = req.ThinkingType cfg.ReasoningEffort = req.ReasoningEffort // API Key 解析优先级：表单输入 > 指定 profile 存的 key > 全局配置的 key。 // 已存 profile 的 key 不回传浏览器，所以测试已存配置时表单为空，需从 DB 取。 @@ -1634,6 +1651,64 @@ func (s *Server) taskScopeList(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"scope": rows}) } +func (s *Server) taskScopeAdd(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + as := s.m.Assets() + if as == nil { + writeErr(w, 503, "asset store 未启用") + return + } + var body struct { + Kind string `json:"kind"` + Value string `json:"value"` + Reason string `json:"reason"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeErr(w, 400, "invalid JSON") + return + } + taskID, _ := strconv.ParseInt(t.ID, 10, 64) + ts, err := as.AddAgentScope(taskID, body.Kind, body.Value, body.Reason, "manual") + if err != nil { + writeErr(w, 400, err.Error()) + return + } + writeJSON(w, 200, ts) +} + +func (s *Server) taskScopeDelete(w http.ResponseWriter, r *http.Request) { + t, ok := s.m.Task(r.PathValue("id")) + if !ok { + writeErr(w, 404, "task not found") + return + } + as := s.m.Assets() + if as == nil { + writeErr(w, 503, "asset store 未启用") + return + } + taskID, _ := strconv.ParseInt(t.ID, 10, 64) + scopeID, err := strconv.ParseInt(r.PathValue("sid"), 10, 64) + if err != nil { + writeErr(w, 400, "invalid scope id") + return + } + deleted, err := as.DeleteTaskScope(taskID, scopeID) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if !deleted { + writeErr(w, 404, "scope row not found") + return + } + writeJSON(w, 200, map[string]bool{"ok": true}) +} + func (s *Server) frontier(w http.ResponseWriter, r *http.Request) { t := s.m.ResolveTask(r.URL.Query().Get("task")) if t == nil { @@ -1833,6 +1908,110 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, dto) } +// findingsExport 导出发现页的漏洞。 +// +// scope = filtered（沿用页面筛选）| all（全部）| selected（勾选的 ids） +// format = md-single（整合一份 .md）| md-zip（一漏洞一 .md,打包 zip） +// | csv | json +// ids = 逗号分隔的 finding id（scope=selected 时必填） +// 筛选参数 severity/status/vulnclass/task_id/sort 与列表接口一致（scope=filtered 用）。 +func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + scope := q.Get("scope") + format := q.Get("format") + + var ids []int64 + var filter db.FindingFilter + switch scope { + case "selected": + for _, part := range strings.Split(q.Get("ids"), ",") { + part = strings.TrimSpace(part) + if part == "" { + continue + } + id, err := strconv.ParseInt(part, 10, 64) + if err != nil || id <= 0 { + writeErr(w, 400, "bad finding id: "+part) + return + } + ids = append(ids, id) + } + if len(ids) == 0 { + writeErr(w, 400, "no findings selected") + return + } + case "all": + // 空 filter = 不加任何条件。 + case "filtered", "": + 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"), + } + default: + writeErr(w, 400, "bad scope: "+scope) + return + } + + fs, err := s.m.pg.ListFindingsForExport(filter, ids) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + + now := time.Now() + stamp := now.Format("20060102-150405") + setDownload := func(contentType, filename string) { + w.Header().Set("Content-Type", contentType) + w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filename)) + } + + switch format { + case "md-single": + setDownload("text/markdown; charset=utf-8", "findings-"+stamp+".md") + _, _ = w.Write([]byte(report.FindingsMarkdown(fs, now))) + case "md-zip": + setDownload("application/zip", "findings-"+stamp+".zip") + zw := zip.NewWriter(w) + used := map[string]int{} + for _, f := range fs { + base := report.FindingFilename(f) + name := base + // 去重:同名文件追加 -2、-3…… + if n := used[base]; n > 0 { + name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1) + } + used[base]++ + fw, werr := zw.Create(name) + if werr != nil { + log.Printf("[findings-export] zip create %s: %v", name, werr) + continue + } + _, _ = fw.Write([]byte(report.SingleFindingMarkdown(f, now))) + } + if cerr := zw.Close(); cerr != nil { + log.Printf("[findings-export] zip close: %v", cerr) + } + case "csv": + setDownload("text/csv; charset=utf-8", "findings-"+stamp+".csv") + _, _ = w.Write(report.FindingsCSV(fs)) + case "json": + assets := s.resolveFindingAssets(fs) + out := make([]FindingDTO, 0, len(fs)) + for _, f := range fs { + out = append(out, findingFromDB(f, assets)) + } + setDownload("application/json; charset=utf-8", "findings-"+stamp+".json") + enc := json.NewEncoder(w) + enc.SetIndent("", " ") + _ = enc.Encode(out) + default: + writeErr(w, 400, "bad format: "+format) + } +} + func findingProvenanceInTask(contextTask *Task, findingTaskID *int64) (sourceTaskID string, inherited, allowed bool) { if contextTask == nil || findingTaskID == nil { return "", false, false @@ -2940,10 +3119,8 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { writeErr(w, 409, "任务正在删除，无法发送新消息") return } - if s.engine.IsPaused(t.ID) { - writeErr(w, 409, "任务已暂停，无法发送新消息") - return - } + // 注意:任务暂停(paused)不拦截主 Agent 对话。主 Agent 编排会话独立于 planner/ + // worker 的暂停,暂停中仍可继续对话(暂停只终止其正在进行的那一轮,见 control())。 var req struct { Message string `json:"message"` Attachments []chatAttachment `json:"attachments,omitempty"` // 方式1 上传的文件(路径相对任务工作目录) @@ -2956,9 +3133,9 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) { // turn busy before its first activity/file write, including rule-mode turns. // If deletion wins the race, the second barrier check rejects this request. s.chatMu.Lock() - if s.engine.IsDeleting(t.ID) || s.engine.IsPaused(t.ID) { + if s.engine.IsDeleting(t.ID) { s.chatMu.Unlock() - writeErr(w, 409, "任务正在删除或已暂停，无法发送新消息") + writeErr(w, 409, "任务正在删除，无法发送新消息") return } if s.chatBusy[t.ID] { diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 9732b9d..8782d56 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -897,6 +897,12 @@ type skillFileNode struct { Compatibility string `json:"compatibility,omitempty"` MCPs []string `json:"mcps,omitempty"` Files []string `json:"files"` + // Usage ledger (db/skill_usage.go). Calls is 0 and LastUsed nil for a skill that + // was never invoked — the ledger only has rows for skills agents actually loaded. + Calls int `json:"calls"` + Tasks int `json:"tasks"` + Agents []string `json:"usage_agents"` + LastUsed *time.Time `json:"last_used,omitempty"` } func (s *Server) fsListSkills(w http.ResponseWriter, r *http.Request) { @@ -915,13 +921,28 @@ func (s *Server) fsListSkills(w http.ResponseWriter, r *http.Request) { writeErr(w, 500, err.Error()) return } + // Usage is best-effort decoration: a ledger read failure leaves the counts at + // zero rather than failing the skill list itself. + statBySkill := map[string]db.SkillStat{} + if s.m.pg != nil { + stats, err := s.m.pg.SkillStats() + if err != nil { + log.Printf("[skills] 读取调用统计失败: %v", err) + } + for _, st := range stats { + statBySkill[st.Skill] = st + } + } nodes := []skillFileNode{} for _, e := range entries { if !e.IsDir() { continue } dirName := e.Name() - node := skillFileNode{Name: dirName} + node := skillFileNode{Name: dirName, Agents: []string{}} + if st, ok := statBySkill[dirName]; ok { + node.Calls, node.Tasks, node.Agents, node.LastUsed = st.Calls, st.Tasks, st.Agents, st.LastUsed + } if meta, ok := metaByDir[dirName]; ok { node.Description = meta.Description node.License = meta.License @@ -937,6 +958,42 @@ func (s *Server) fsListSkills(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"skills": nodes}) } +// fsSkillUsage returns one skill's recent invocations, newest first (detail panel). +func (s *Server) fsSkillUsage(w http.ResponseWriter, r *http.Request) { + pg := s.pg(w) + if pg == nil { + return + } + name := r.PathValue("name") + if !validSkillName(name) { + writeErr(w, 400, "非法 skill 名") + return + } + limit, _ := strconv.Atoi(r.URL.Query().Get("limit")) + calls, err := pg.RecentSkillCalls(name, limit) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"calls": calls}) +} + +// fsMissingSkills lists skill names agents asked for that do not exist — the gap +// list, i.e. which procedures are worth writing next. +func (s *Server) fsMissingSkills(w http.ResponseWriter, r *http.Request) { + pg := s.pg(w) + if pg == nil { + return + } + limit, _ := strconv.Atoi(r.URL.Query().Get("limit")) + missing, err := pg.MissingSkillStats(limit) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"missing": missing}) +} + // cleanStrs trims each string and drops empties. func cleanStrs(in []string) []string { out := make([]string, 0, len(in)) diff --git a/server/skill_usage.go b/server/skill_usage.go new file mode 100644 index 0000000..4d0906e --- /dev/null +++ b/server/skill_usage.go @@ -0,0 +1,88 @@ +package server + +import ( + "context" + "encoding/json" + "log" + "path/filepath" + "strings" + + "github.com/Autumn-27/norma/skill" + actool "github.com/Autumn-27/norma/tool" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" +) + +// maxLedgerSkillName caps the skill name stored for a MISS — that string comes +// straight from the model and is otherwise unbounded. +const maxLedgerSkillName = 128 + +// meteredSkill wraps the Skill meta-tool to append one skill_usage row per call. +// The metering lives here rather than in Registry.OnInvoke because OnInvoke only +// receives the resolved Skill: it never fires for an unknown skill name (norma +// returns a tool error before the hook) and it cannot see the caller's args. Both +// matter — a miss tells you which procedure the agents wished existed. +type meteredSkill struct { + actool.CoreTool + pg *db.DB + reg *skill.Registry + agentKey string + ri agent.RunInfo +} + +// meterSkillTool returns t unchanged when there is no DB to meter into. +func meterSkillTool(t actool.CoreTool, pg *db.DB, reg *skill.Registry, agentKey string, ri agent.RunInfo) actool.CoreTool { + if pg == nil { + return t + } + return &meteredSkill{CoreTool: t, pg: pg, reg: reg, agentKey: agentKey, ri: ri} +} + +func (m *meteredSkill) Call(ctx context.Context, input json.RawMessage, tc *actool.ToolContext) (actool.Result, error) { + m.record(input) + return m.CoreTool.Call(ctx, input, tc) +} + +// record appends the ledger row. Best-effort by design: a metering failure must +// never turn into a failed skill invocation, so every error path just logs. +func (m *meteredSkill) record(input json.RawMessage) { + var in struct { + Name string `json:"name"` + Args string `json:"args"` + } + if err := json.Unmarshal(input, &in); err != nil { + return // malformed call; norma will reject it too + } + name := strings.TrimSpace(in.Name) + if name == "" { + return + } + // Resolve to the directory name so the ledger keys match agent_skill_visibility + // and the skills page (a skill's display Name can differ from its directory). + found := false + if m.reg != nil { + if s, ok := m.reg.Get(name); ok { + found = true + if s.Dir != "" { + name = filepath.Base(s.Dir) + } + } + } + if len(name) > maxLedgerSkillName { + name = name[:maxLedgerSkillName] + } + err := m.pg.InsertSkillUsage(&db.SkillUsage{ + Skill: name, + AgentKey: m.agentKey, + TaskID: m.ri.TaskID, + ExplorationID: m.ri.ExplorationID, + IntentID: m.ri.IntentID, + SessionID: m.ri.SessionID, + ArgsLen: len(in.Args), + Found: found, + }) + if err != nil { + log.Printf("[skillusage] insert: %v", err) + } +} diff --git a/web/src/app/(main)/chat/page.tsx b/web/src/app/(main)/chat/page.tsx index 2b6a7fb..32ecdb2 100644 --- a/web/src/app/(main)/chat/page.tsx +++ b/web/src/app/(main)/chat/page.tsx @@ -77,8 +77,8 @@ function LiveBadge() { ); } -// Composer is the shared bottom input (textarea grows to a cap, Enter newlines, -// Ctrl/Cmd+Enter sends) — the same affordance across DraftChat and ChatView. +// Composer is the shared bottom input (textarea grows to a cap, Enter sends, +// Shift+Enter newlines) — the same affordance across DraftChat and ChatView. function Composer({ value, onChange, @@ -112,11 +112,10 @@ function Composer({ const fileInputRef = React.useRef(null); const atts = attachments ?? []; function onKeyDown(e: React.KeyboardEvent) { - // Enter 换行；只有 Ctrl/Cmd+Enter 才发送(避免长消息误发)。 - if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) { - e.preventDefault(); - onSend(); - } + // Enter 发送；Shift+Enter 换行。isComposing 时是中文等输入法在选字,不拦截。 + if (e.key !== "Enter" || e.shiftKey || e.nativeEvent.isComposing) return; + e.preventDefault(); + onSend(); } return (
@@ -364,7 +363,7 @@ function DraftChat({ onChange={setInput} onSend={send} disabled={sending || uploading || !agentKey} - placeholder="输入消息，Enter 换行，Ctrl+Enter 发送" + placeholder="输入消息，Enter 发送，Shift+Enter 换行" leftSlot={agentPicker} onPickFiles={pickFiles} uploading={uploading} @@ -692,7 +691,7 @@ function ChatView({ onChange={setInput} onSend={send} disabled={running} - placeholder={running ? "Agent 正在回复…" : "输入消息，Enter 换行，Ctrl+Enter 发送"} + placeholder={running ? "Agent 正在回复…" : "输入消息，Enter 发送，Shift+Enter 换行"} running={running} onStop={stop} stopDisabled={stopping} diff --git a/web/src/app/(main)/dashboard/page.tsx b/web/src/app/(main)/dashboard/page.tsx index 9aaab90..881c0b9 100644 --- a/web/src/app/(main)/dashboard/page.tsx +++ b/web/src/app/(main)/dashboard/page.tsx @@ -1,7 +1,9 @@ "use client"; import * as React from "react"; + import Link from "next/link"; + import { ActivityIcon, ArrowUpRightIcon, @@ -13,29 +15,18 @@ import { TargetIcon, ZapIcon, } from "lucide-react"; -import { - BarChart, - Bar, - XAxis, - YAxis, - CartesianGrid, - Tooltip as RechartsTooltip, -} from "recharts"; -import { - ChartContainer, - ChartTooltip, - ChartTooltipContent, - type ChartConfig, -} from "@/components/ui/chart"; -import { cn } from "@/lib/utils"; +import { Bar, BarChart, CartesianGrid, Tooltip as RechartsTooltip, XAxis, YAxis } from "recharts"; + +import { StatusBadge } from "@/components/status-badge"; import { Badge } from "@/components/ui/badge"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { type ChartConfig, ChartContainer, ChartTooltip, ChartTooltipContent } from "@/components/ui/chart"; import { Progress } from "@/components/ui/progress"; -import { StatusBadge } from "@/components/status-badge"; import { api } from "@/lib/api"; import type { Activity, Agent, + ConvTokenSummary, Finding, InterceptPending, LLMProfile, @@ -47,14 +38,15 @@ import type { TokenTotal, Tool, TrafficExchange, - ConvTokenSummary, + UsageStats, } from "@/lib/types"; +import { cn } from "@/lib/utils"; // ── chart constants ─────────────────────────────────────────────────────────── const dailyTrendConfig = { - input: { label: "输入", color: "hsl(217 91% 60%)" }, - output: { label: "输出", color: "hsl(263 70% 60%)" }, + input: { label: "输入", color: "hsl(217 91% 60%)" }, + output: { label: "输出", color: "hsl(263 70% 60%)" }, cacheRead: { label: "缓存读", color: "hsl(160 60% 45%)" }, } satisfies ChartConfig; @@ -62,9 +54,7 @@ const dailyTrendConfig = { function fmtRel(ts?: string | number): string { if (!ts) return "—"; - const ms = - Date.now() - - (typeof ts === "number" ? ts * 1000 : Date.parse(ts as string)); + const ms = Date.now() - (typeof ts === "number" ? ts * 1000 : Date.parse(ts as string)); const s = Math.floor(ms / 1000); if (s < 60) return `${s}s 前`; const m = Math.floor(s / 60); @@ -115,14 +105,7 @@ function statusBg(code: number): string { // ── sub-components ──────────────────────────────────────────────────────────── function LiveDot({ className }: { className?: string }) { - return ( - - ); + return ; } function SectionTitle({ @@ -157,6 +140,7 @@ export default function DashboardPage() { const [activity, setActivity] = React.useState([]); const [tokens, setTokens] = React.useState(null); const [convTokens, setConvTokens] = React.useState([]); + const [usageStats, setUsageStats] = React.useState(null); const [assetCounts, setAssetCounts] = React.useState>({}); const [traffic, setTraffic] = React.useState([]); const [agents, setAgents] = React.useState([]); @@ -195,7 +179,10 @@ export default function DashboardPage() { }; load(); const t = setInterval(load, 5000); - return () => { alive = false; clearInterval(t); }; + return () => { + alive = false; + clearInterval(t); + }; }, []); // slow poll: traffic, assets, system-static (every 15s) @@ -203,16 +190,16 @@ export default function DashboardPage() { let alive = true; const load = async () => { try { - const [traf, counts, agentList, mcpList, skillList, toolList, profileList] = - await Promise.all([ - api.traffic(0, 50), - api.assetCounts(), - api.agents(), - api.mcpServers(), - api.skills(), - api.tools(), - api.llmProfiles(), - ]); + const [traf, counts, agentList, mcpList, skillList, toolList, profileList, usage] = await Promise.all([ + api.traffic(0, 50), + api.assetCounts(), + api.agents(), + api.mcpServers(), + api.skills(), + api.tools(), + api.llmProfiles(), + api.usageStats(365), + ]); if (!alive) return; setTraffic(traf.exchanges ?? []); setAssetCounts(counts ?? {}); @@ -221,13 +208,17 @@ export default function DashboardPage() { setSkills(skillList); setTools(toolList); setLLMProfiles(profileList); + setUsageStats(usage); } catch { /* transient errors */ } }; load(); const t = setInterval(load, 15000); - return () => { alive = false; clearInterval(t); }; + return () => { + alive = false; + clearInterval(t); + }; }, []); // ── derived ─────────────────────────────────────────────────────────────── @@ -249,20 +240,13 @@ export default function DashboardPage() { const sortedTasks = React.useMemo( () => [...tasks] - .sort( - (a, b) => - (b.last_activity_unix ?? b.created_unix ?? 0) - - (a.last_activity_unix ?? a.created_unix ?? 0), - ) + .sort((a, b) => (b.last_activity_unix ?? b.created_unix ?? 0) - (a.last_activity_unix ?? a.created_unix ?? 0)) .slice(0, 5), [tasks], ); const recentFindings = React.useMemo( - () => - [...findings] - .sort((a, b) => Date.parse(b.ts) - Date.parse(a.ts)) - .slice(0, 8), + () => [...findings].sort((a, b) => Date.parse(b.ts) - Date.parse(a.ts)).slice(0, 8), [findings], ); @@ -275,10 +259,7 @@ export default function DashboardPage() { [activity], ); - const totalAssets = React.useMemo( - () => Object.values(assetCounts).reduce((a, b) => a + b, 0), - [assetCounts], - ); + const totalAssets = React.useMemo(() => Object.values(assetCounts).reduce((a, b) => a + b, 0), [assetCounts]); // asset type breakdown const assetByType = React.useMemo(() => { @@ -303,9 +284,7 @@ export default function DashboardPage() { }, [traffic]); const trafficMax = Math.max(...trafficByCodes.map((x) => x.n), 1); - const recentTraffic = [...traffic] - .sort((a, b) => Date.parse(b.ts) - Date.parse(a.ts)) - .slice(0, 5); + const recentTraffic = [...traffic].sort((a, b) => Date.parse(b.ts) - Date.parse(a.ts)).slice(0, 5); // system const activeProfile = llmProfiles.find((p) => p.is_default); @@ -316,100 +295,147 @@ export default function DashboardPage() { // tasks whose llm_profile_id is null/undefined used the active default profile const defaultProfileId = activeProfile ? Number(activeProfile.id) : null; - const tokenByProfile = React.useMemo< - Map - >(() => { - const m = new Map(); - for (const t of tasks) { - // null means "used whatever was active" → bucket under defaultProfileId - const key = t.llm_profile_id != null ? t.llm_profile_id : defaultProfileId; - const prev = m.get(key) ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; - m.set(key, { - input: prev.input + (t.tokens?.input_tokens ?? 0), - output: prev.output + (t.tokens?.output_tokens ?? 0), - cacheRead: prev.cacheRead + (t.tokens?.cache_read_tokens ?? 0), - cacheWrite: prev.cacheWrite + (t.tokens?.cache_write_tokens ?? 0), - taskCount: prev.taskCount + 1, - }); - } - // merge conversation (chat) usage into the same per-profile buckets — token sums - // only; taskCount stays a task count. - for (const c of convTokens) { - const key = c.llm_profile_id != null ? c.llm_profile_id : defaultProfileId; - const prev = m.get(key) ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; - m.set(key, { - input: prev.input + c.input_tokens, - output: prev.output + c.output_tokens, - cacheRead: prev.cacheRead + c.cache_read_tokens, - cacheWrite: prev.cacheWrite + c.cache_write_tokens, - taskCount: prev.taskCount, - }); - } - return m; - }, [tasks, convTokens, defaultProfileId]); - + // 数据源开关：旧版 = activity（task.tokens + 会话），新版 = llm_usage 计量账本。 + const [tokenVersion, setTokenVersion] = React.useState<"old" | "new">("old"); // selected profile tab: "all" = 全部; number = specific profile id const [tokenTab, setTokenTab] = React.useState("all"); // day range for the daily bar chart const [tokenDays, setTokenDays] = React.useState<7 | 30 | 90 | 180 | 365>(30); + // profile 名 → id，用于把 llm_usage 的 profile_name 映射到现有 profile 分栏。 + const profileIdByName = React.useMemo(() => { + const m = new Map(); + for (const p of llmProfiles) m.set(p.name, Number(p.id)); + return m; + }, [llmProfiles]); + + type Bucket = { input: number; output: number; cacheRead: number; cacheWrite: number; taskCount: number }; + + // 旧版：按 profile 归桶（来自 activity 的 task.tokens + 会话用量）。 + const tokenByProfileOld = React.useMemo>(() => { + const m = new Map(); + const fold = (key: number | null, inp: number, out: number, cr: number, cw: number, addTask: boolean) => { + const prev = m.get(key) ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; + m.set(key, { + input: prev.input + inp, + output: prev.output + out, + cacheRead: prev.cacheRead + cr, + cacheWrite: prev.cacheWrite + cw, + taskCount: prev.taskCount + (addTask ? 1 : 0), + }); + }; + for (const t of tasks) { + fold( + t.llm_profile_id ?? defaultProfileId, + t.tokens?.input_tokens ?? 0, + t.tokens?.output_tokens ?? 0, + t.tokens?.cache_read_tokens ?? 0, + t.tokens?.cache_write_tokens ?? 0, + true, + ); + } + for (const c of convTokens) { + fold( + c.llm_profile_id ?? defaultProfileId, + c.input_tokens, + c.output_tokens, + c.cache_read_tokens, + c.cache_write_tokens, + false, + ); + } + return m; + }, [tasks, convTokens, defaultProfileId]); + + // 新版：按 profile 归桶（来自 llm_usage 全局聚合，逐次调用精确）。 + const tokenByProfileNew = React.useMemo>(() => { + const m = new Map(); + for (const p of usageStats?.by_profile ?? []) { + // 未匹配到现有 profile（改名/删除/空名）→ 落到默认桶，仍计入「全部」。 + const key = profileIdByName.get(p.profile_name) ?? defaultProfileId; + const prev = m.get(key) ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; + m.set(key, { + input: prev.input + p.input_tokens, + output: prev.output + p.output_tokens, + cacheRead: prev.cacheRead + p.cache_read_tokens, + cacheWrite: prev.cacheWrite + p.cache_write_tokens, + taskCount: prev.taskCount + p.tasks, + }); + } + return m; + }, [usageStats, profileIdByName, defaultProfileId]); + + const tokenByProfile = tokenVersion === "new" ? tokenByProfileNew : tokenByProfileOld; + const displayedTokens = React.useMemo(() => { if (tokenTab === "all") { - // sum across all profiles - let input = 0, output = 0, cacheRead = 0, cacheWrite = 0, taskCount = 0; + let input = 0, + output = 0, + cacheRead = 0, + cacheWrite = 0, + taskCount = 0; for (const v of tokenByProfile.values()) { - input += v.input; output += v.output; - cacheRead += v.cacheRead; cacheWrite += v.cacheWrite; + input += v.input; + output += v.output; + cacheRead += v.cacheRead; + cacheWrite += v.cacheWrite; taskCount += v.taskCount; } return { input, output, cacheRead, cacheWrite, taskCount }; } const v = tokenByProfile.get(tokenTab as number | null); - return v ? { input: v.input, output: v.output, cacheRead: v.cacheRead, cacheWrite: v.cacheWrite, taskCount: v.taskCount } - : { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; + return v + ? { input: v.input, output: v.output, cacheRead: v.cacheRead, cacheWrite: v.cacheWrite, taskCount: v.taskCount } + : { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, taskCount: 0 }; }, [tokenTab, tokenByProfile]); - // daily token data: tasks grouped by date, filtered by day range and selected profile - const dailyTokenData = React.useMemo(() => { + // 旧版每日：把任务/会话的总量按其创建日期归桶（近似，非真实每日消耗）。 + const dailyTokenDataOld = React.useMemo(() => { const now = new Date(); now.setDate(now.getDate() - tokenDays); const cutoff = now.toISOString().slice(0, 10); - const m = new Map(); + const fold = (date: string, inp: number, out: number, cr: number) => { + if (date < cutoff) return; + const prev = m.get(date) ?? { input: 0, output: 0, cacheRead: 0 }; + m.set(date, { input: prev.input + inp, output: prev.output + out, cacheRead: prev.cacheRead + cr }); + }; for (const t of tasks) { - // filter by profile tab when not "all" - if (tokenTab !== "all") { - const taskKey = t.llm_profile_id ?? defaultProfileId; - if (taskKey !== tokenTab) continue; - } - const date = t.created_at.slice(0, 10); - if (date < cutoff) continue; - const prev = m.get(date) ?? { input: 0, output: 0, cacheRead: 0 }; - m.set(date, { - input: prev.input + (t.tokens?.input_tokens ?? 0), - output: prev.output + (t.tokens?.output_tokens ?? 0), - cacheRead: prev.cacheRead + (t.tokens?.cache_read_tokens ?? 0), - }); + if (tokenTab !== "all" && (t.llm_profile_id ?? defaultProfileId) !== tokenTab) continue; + fold( + t.created_at.slice(0, 10), + t.tokens?.input_tokens ?? 0, + t.tokens?.output_tokens ?? 0, + t.tokens?.cache_read_tokens ?? 0, + ); } - // merge conversation usage by conversation created_at date, same profile filter. for (const c of convTokens) { - if (tokenTab !== "all") { - const convKey = c.llm_profile_id ?? defaultProfileId; - if (convKey !== tokenTab) continue; - } - const date = c.created_at.slice(0, 10); - if (date < cutoff) continue; - const prev = m.get(date) ?? { input: 0, output: 0, cacheRead: 0 }; - m.set(date, { - input: prev.input + c.input_tokens, - output: prev.output + c.output_tokens, - cacheRead: prev.cacheRead + c.cache_read_tokens, + if (tokenTab !== "all" && (c.llm_profile_id ?? defaultProfileId) !== tokenTab) continue; + fold(c.created_at.slice(0, 10), c.input_tokens, c.output_tokens, c.cache_read_tokens); + } + return [...m.entries()].sort(([a], [b]) => a.localeCompare(b)).map(([date, v]) => ({ date: date.slice(5), ...v })); + }, [tasks, convTokens, tokenDays, tokenTab, defaultProfileId]); + + // 新版每日：来自 llm_usage 的真实每日消耗（ts 是实际调用时刻）。 + const dailyTokenDataNew = React.useMemo(() => { + const now = new Date(); + now.setDate(now.getDate() - tokenDays); + const cutoff = now.toISOString().slice(0, 10); + const m = new Map(); + for (const d of usageStats?.daily ?? []) { + if (d.date < cutoff) continue; + if (tokenTab !== "all" && (profileIdByName.get(d.profile_name) ?? defaultProfileId) !== tokenTab) continue; + const prev = m.get(d.date) ?? { input: 0, output: 0, cacheRead: 0 }; + m.set(d.date, { + input: prev.input + d.input_tokens, + output: prev.output + d.output_tokens, + cacheRead: prev.cacheRead + d.cache_read_tokens, }); } - return [...m.entries()] - .sort(([a], [b]) => a.localeCompare(b)) - .map(([date, v]) => ({ date: date.slice(5), ...v })); - }, [tasks, convTokens, tokenDays, tokenTab, defaultProfileId]); + return [...m.entries()].sort(([a], [b]) => a.localeCompare(b)).map(([date, v]) => ({ date: date.slice(5), ...v })); + }, [usageStats, tokenDays, tokenTab, profileIdByName, defaultProfileId]); + + const dailyTokenData = tokenVersion === "new" ? dailyTokenDataNew : dailyTokenDataOld; // activity kind label function kindLabel(a: Activity): string { @@ -434,7 +460,6 @@ export default function DashboardPage() { return (
- {/* ── Header ── */}
@@ -451,7 +476,6 @@ export default function DashboardPage() { {/* ── Row 1: 5 stat cards ── */}
- {/* 活跃任务 */} @@ -459,25 +483,15 @@ export default function DashboardPage() { 活跃任务
- - {tasksByStatus.running ?? 0} - + {tasksByStatus.running ?? 0} / {tasks.length}
- {(tasksByStatus.running ?? 0) > 0 && ( - 探索 {tasksByStatus.running} - )} - {(tasksByStatus.paused ?? 0) > 0 && ( - 暂停 {tasksByStatus.paused} - )} - {(tasksByStatus.done ?? 0) > 0 && ( - 完成 {tasksByStatus.done} - )} - {tasks.length === 0 && ( - 暂无任务 - )} + {(tasksByStatus.running ?? 0) > 0 && 探索 {tasksByStatus.running}} + {(tasksByStatus.paused ?? 0) > 0 && 暂停 {tasksByStatus.paused}} + {(tasksByStatus.done ?? 0) > 0 && 完成 {tasksByStatus.done}} + {tasks.length === 0 && 暂无任务} @@ -487,9 +501,7 @@ export default function DashboardPage() {
确认发现
-
- {findings.length} -
+
{findings.length}
严重 {findingsBySev.critical} @@ -505,13 +517,9 @@ export default function DashboardPage() {
资产节点
-
- {totalAssets} -
+
{totalAssets}
- - 跨任务共享 - + 跨任务共享 {/* 流量交互 */} @@ -520,9 +528,7 @@ export default function DashboardPage() {
流量交互
-
- {traffic.length} -
+
{traffic.length}
{settings?.traffic_capture ? ( @@ -547,7 +553,8 @@ export default function DashboardPage() {
- 入 {fmtTokens(displayedTokens.input)}（含缓存 {fmtTokens(displayedTokens.cacheRead)}）· 出 {fmtTokens(displayedTokens.output)} + 入 {fmtTokens(displayedTokens.input)}（含缓存 {fmtTokens(displayedTokens.cacheRead)}）· 出{" "} + {fmtTokens(displayedTokens.output)}
@@ -559,8 +566,35 @@ export default function DashboardPage() {
LLM Token 消耗 - - {tasks.length} 个任务累计 + {/* 数据源开关：旧版=activity 统计（含历史任务），新版=llm_usage 计量账本（更准，仅覆盖启用后） */} +
+ {( + [ + { v: "old", label: "旧版" }, + { v: "new", label: "新版" }, + ] as const + ).map(({ v, label }) => ( + + ))} +
+ + {tokenVersion === "new" ? "llm_usage" : "activity"}
{/* Profile tabs */} @@ -579,7 +613,7 @@ export default function DashboardPage() { {llmProfiles.map((p) => { const key = Number(p.id); const hasData = tokenByProfile.has(key) || (p.is_default && tokenByProfile.has(defaultProfileId)); - const resolvedKey = tokenByProfile.has(key) ? key : (p.is_default ? defaultProfileId : key); + const resolvedKey = tokenByProfile.has(key) ? key : p.is_default ? defaultProfileId : key; return (
{dailyTokenData.length === 0 ? ( -
+
暂无数据
) : ( - + (
- {(dailyTrendConfig as Record)[String(name)]?.label ?? String(name)} - - {fmtTokens(Number(value))} + + {(dailyTrendConfig as Record)[String(name)]?.label ?? + String(name)} + {fmtTokens(Number(value))}
)} /> } /> - - + + @@ -765,7 +815,6 @@ export default function DashboardPage() { {/* ── Row 3: 活动流 | 发现 ── */}
- {/* 活动流 */}
@@ -787,27 +836,20 @@ export default function DashboardPage() {
- {(["tool_use", "tool_result", "text", "thinking", "result"] as const).map( - (kind) => { - const count = activity.filter((a) => a.kind === kind).length; - if (count === 0) return null; - return ( - - {kind} {count} - - ); - }, - )} + {(["tool_use", "tool_result", "text", "thinking", "result"] as const).map((kind) => { + const count = activity.filter((a) => a.kind === kind).length; + if (count === 0) return null; + return ( + + {kind} {count} + + ); + })}
{recentActivity.length === 0 ? ( -
- 暂无活动记录 -
+
暂无活动记录
) : ( recentActivity.map((a) => (
@@ -821,20 +863,13 @@ export default function DashboardPage() {
{kindLabel(a)}
-
- {a.summary} -
+
{a.summary}
- + {a.kind} -
- {fmtRel(a.ts)} -
+
{fmtRel(a.ts)}
)) @@ -859,9 +894,7 @@ export default function DashboardPage() {
{recentFindings.length === 0 ? ( -
- 暂无发现 -
+
暂无发现
) : ( recentFindings.map((f) => (
@@ -874,18 +907,14 @@ export default function DashboardPage() {
{f.vulnclass} {f.task_description && ( - - {f.task_description} - + {f.task_description} )}
{f.summary}
- - {fmtRel(f.ts)} - + {fmtRel(f.ts)}
)) )} @@ -901,9 +930,7 @@ export default function DashboardPage() { 任务
- - {tasks.length} 个任务 - + {tasks.length} 个任务 - {["任务", "状态", "引擎", "目标进度", "在途", "最近活动"].map( - (h) => ( - - {h} - - ), - )} + {["任务", "状态", "引擎", "目标进度", "在途", "最近活动"].map((h) => ( + + {h} + + ))} {sortedTasks.length === 0 ? ( - + 暂无任务 ) : ( sortedTasks.map((t) => { - const goalsPct = t.goals_total - ? Math.round(((t.goals_met ?? 0) / t.goals_total) * 100) - : null; + const goalsPct = t.goals_total ? Math.round(((t.goals_met ?? 0) / t.goals_total) * 100) : null; return ( - + - - - {t.description} - - - {t.id} - + + {t.description} + {t.id} @@ -983,7 +993,7 @@ export default function DashboardPage() { )} - {t.in_flight ?? 0 > 0 ? ( + {(t.in_flight ?? 0 > 0) ? ( {t.in_flight} ) : ( — @@ -1002,7 +1012,6 @@ export default function DashboardPage() { {/* ── Row 5: 资产分布 | 流量状态码 | 拦截 & 待审批 ── */}
- {/* 资产分布 */} @@ -1010,9 +1019,7 @@ export default function DashboardPage() { {assetByType.length === 0 ? ( -
- 暂无资产数据 -
+
暂无资产数据
) : (
{assetByType.map(([type, count]) => ( @@ -1032,9 +1039,7 @@ export default function DashboardPage() {
)} -
- 共 {totalAssets} 节点 -
+
共 {totalAssets} 节点
 {/* 流量状态码 */} @@ -1045,36 +1050,24 @@ export default function DashboardPage() { {/* bar chart */} {trafficByCodes.length === 0 ? ( -
- 暂无流量数据 -
+
暂无流量数据
) : ( <>
{trafficByCodes.map(({ code, n }) => (
- - {n} - + {n}
- - {code} - + {code}
))}
-
- 最近请求 -
+
最近请求
{recentTraffic.map((e) => (
@@ -1092,14 +1085,7 @@ export default function DashboardPage() { {e.host} {e.url.replace(/^https?:\/\/[^/]+/, "").substring(0, 30)} - - {e.status} - + {e.status}
))}
@@ -1154,9 +1140,7 @@ export default function DashboardPage() { {/* pending approvals */} {pendingCount > 0 && (
-
- 待审批 ({pendingCount}) -
+
待审批 ({pendingCount})
{pending.slice(0, 3).map((p) => (
{p.tool_name}
-
- {p.agent_name} -
+
{p.agent_name}
diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index d5590e0..7d36f5c 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -20,6 +20,7 @@ import { SelectTrigger, } from "@/components/ui/select"; import { StatusBadge } from "@/components/status-badge"; +import { CopyButton } from "@/components/copy-button"; import { Markdown } from "@/components/markdown"; import { FindingLineageView } from "./lineage"; import { statusMeta } from "@/lib/status"; @@ -208,8 +209,11 @@ function FindingDetailInner() { {/* 证据下方：详细报告(Markdown 渲染) */} - + 详细报告 + {finding.report && ( + + )} {finding.report ? ( diff --git a/web/src/app/(main)/function/findings/page.tsx b/web/src/app/(main)/function/findings/page.tsx index 38b9eda..03201b3 100644 --- a/web/src/app/(main)/function/findings/page.tsx +++ b/web/src/app/(main)/function/findings/page.tsx @@ -8,10 +8,22 @@ import { ArrowUpRightIcon, FileTextIcon, Trash2Icon, + DownloadIcon, } from "lucide-react"; import { StatusBadge } from "@/components/status-badge"; +import { CopyButton } from "@/components/copy-button"; import { Markdown } from "@/components/markdown"; +import { Checkbox } from "@/components/ui/checkbox"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { AlertDialog, AlertDialogAction, @@ -101,6 +113,73 @@ export default function FindingsPage() { const [page, setPage] = React.useState(1); const [pageSize, setPageSize] = React.useState(20); + // 勾选导出:按 finding_id(独立表 id)记选中项,跨页保留。 + const [selectedIds, setSelectedIds] = React.useState>(() => new Set()); + // 导出弹窗状态:范围(当前筛选/全部/选中) × 格式(md 单文件/md 分文件 zip/csv/json)。 + const [exportOpen, setExportOpen] = React.useState(false); + const [exportScope, setExportScope] = React.useState<"filtered" | "all" | "selected">("filtered"); + const [exportFormat, setExportFormat] = React.useState<"md-single" | "md-zip" | "csv" | "json">("md-single"); + const [exporting, setExporting] = React.useState(false); + + const toggleSelected = React.useCallback((id: string, checked: boolean) => { + setSelectedIds((prev) => { + const next = new Set(prev); + if (checked) next.add(id); + else next.delete(id); + return next; + }); + }, []); + + // 当前页可勾选的 finding_id(排除无 finding_id 的继承行)。 + const pageSelectableIds = React.useMemo( + () => findings.map((f) => f.finding_id).filter((id): id is string => !!id), + [findings], + ); + const pageSelectedCount = pageSelectableIds.filter((id) => selectedIds.has(id)).length; + let headerChecked: boolean | "indeterminate" = false; + if (pageSelectableIds.length > 0 && pageSelectedCount === pageSelectableIds.length) { + headerChecked = true; + } else if (pageSelectedCount > 0) { + headerChecked = "indeterminate"; + } + const toggleSelectedPage = React.useCallback( + (checked: boolean) => { + setSelectedIds((prev) => { + const next = new Set(prev); + for (const id of pageSelectableIds) { + if (checked) next.add(id); + else next.delete(id); + } + return next; + }); + }, + [pageSelectableIds], + ); + + // 打开导出弹窗时,若有勾选项则默认范围切到「选中」,否则「当前筛选」。 + function openExport() { + setExportScope(selectedIds.size > 0 ? "selected" : "filtered"); + setExportOpen(true); + } + + async function doExport() { + setExporting(true); + try { + await api.exportFindings({ + format: exportFormat, + scope: exportScope, + filters: { severity, status, vulnclass, task, sort }, + ids: [...selectedIds], + }); + setExportOpen(false); + toast.success("已开始下载导出文件"); + } catch (e) { + toast.error("导出失败：" + (e as Error).message); + } finally { + setExporting(false); + } + } + // reset to page 1 whenever filters change React.useEffect(() => { setPage(1); @@ -386,9 +465,19 @@ export default function FindingsPage() { - - 共 {total} 条 - +
+ {selectedIds.size > 0 && ( + + 已选 {selectedIds.size} 条 + + )} + + 共 {total} 条 + + +
@@ -398,6 +487,13 @@ export default function FindingsPage() { + + toggleSelectedPage(c === true)} + aria-label="选择本页全部" + /> + 严重度 漏洞名称 @@ -417,6 +513,17 @@ export default function FindingsPage() { className="cursor-pointer" onClick={() => toggleRow(f)} > + e.stopPropagation()}> + {f.finding_id && ( + + toggleSelected(f.finding_id as string, c === true) + } + aria-label="选择该漏洞" + /> + )} + {/* whitespace-normal 覆盖 TableCell 默认的 nowrap,否则展开区文字 被强制单行、直接溢出单元格。 */} - +
{/* 行内编辑:名称/类别/严重等级,可改并保存(仅独立 finding 行)。 */} {f.finding_id && edit && ( @@ -643,9 +750,20 @@ export default function FindingsPage() { {/* 详细报告(Markdown):展开时按 finding_id 懒加载,免进详情页即可查看。 */} {f.finding_id && (
-
- - 详细报告 +
+ + + 详细报告 + + {reports[f.id]?.status === "done" && + reports[f.id]?.text.trim() && ( + + )}
{(() => { const rep = reports[f.id]; @@ -686,7 +804,7 @@ export default function FindingsPage() { {findings.length === 0 && ( 没有匹配的发现。 @@ -705,6 +823,76 @@ export default function FindingsPage() {
+ + + + + 导出发现 + + 选择导出范围与格式,生成后浏览器会自动下载。 + + + +
+
+ + setExportScope(v as typeof exportScope)} + > + + + + +
+ +
+ + setExportFormat(v as typeof exportFormat)} + > + + + + + +
+
+ + + + + + +
); } diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx index f5fd9de..3c06a2b 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/overview-tab.tsx @@ -1,27 +1,63 @@ "use client"; import * as React from "react"; + import { ActivityIcon, AlertTriangleIcon, BugIcon, ClockIcon, + CoinsIcon, + PlusIcon, RefreshCwIcon, ShieldCheckIcon, TargetIcon, + Trash2Icon, } from "lucide-react"; -import { - Card, - CardContent, - CardDescription, - CardHeader, - CardTitle, -} from "@/components/ui/card"; -import { Button } from "@/components/ui/button"; -import { Progress } from "@/components/ui/progress"; + import { StatusBadge } from "@/components/status-badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { NativeSelect, NativeSelectOption } from "@/components/ui/native-select"; +import { Progress } from "@/components/ui/progress"; import { api } from "@/lib/api"; -import type { Task, Stats, TaskNode, Finding } from "@/lib/types"; +import type { Finding, ModelTokenStat, Stats, Task, TaskNode, TaskScopeRow } from "@/lib/types"; + +// 紧凑格式化 token 数（12345 → 12.3k，2000000 → 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"; + return String(n); +} + +// 缓存命中率 = 缓存读 / 输入（InputTokens 已含 cache_read 子集，故比值在 0–100%）。 +function cacheHitRate(cacheRead: number, input: number): string { + if (input <= 0) return "—"; + return Math.round((cacheRead / input) * 100) + "%"; +} + +// 测试范围一条的显示值：域名 / 网段 / 公司。 +function scopeValue(row: TaskScopeRow): string { + if (row.domain) return row.domain; + if (row.net) return row.net; + if (row.company_id) return `company #${row.company_id}`; + return "—"; +} + +const SCOPE_KIND_LABELS: Record = { + company: "公司", + root_domain: "根域名", + subdomain: "子域名", + ip: "IP", + cidr: "网段", +}; + +const SCOPE_SOURCE_LABELS: Record = { + auto: "自动", + agent: "Agent", + manual: "手动", +}; function StatCard({ label, @@ -42,9 +78,7 @@ function StatCard({ {value} - {sub && ( - {sub} - )} + {sub && {sub}} ); } @@ -63,6 +97,57 @@ export function OverviewTab({ taskId }: { taskId: string }) { } | null>(null); // 正在重跑的意图 id（含 "__all__" 表示批量），用于禁用按钮 + 转圈。 const [rerunning, setRerunning] = React.useState>(new Set()); + // 测试范围列表 + 新增表单状态。 + const [scope, setScope] = React.useState([]); + const [scopeKind, setScopeKind] = React.useState("root_domain"); + const [scopeValueInput, setScopeValueInput] = React.useState(""); + const [scopeBusy, setScopeBusy] = React.useState(false); + const [scopeErr, setScopeErr] = React.useState(""); + // 按模型的 token 用量（来自常开的 llm_usage 计量账本，逐次精确）。 + const [modelTokens, setModelTokens] = React.useState([]); + + const loadTokens = React.useCallback(async () => { + try { + const resp = await api.tokensByModel(taskId); + setModelTokens(resp.models); + } catch { + // 忽略：无 PG 时接口报错，卡片自然为空 + } + }, [taskId]); + + const loadScope = React.useCallback(async () => { + try { + const resp = await api.taskScope(taskId); + setScope(resp.scope); + } catch { + // 忽略：无 asset store 时接口 503，范围卡片自然为空 + } + }, [taskId]); + + const addScope = async () => { + const value = scopeValueInput.trim(); + if (!value) return; + setScopeBusy(true); + setScopeErr(""); + try { + await api.addTaskScope(taskId, scopeKind, value); + setScopeValueInput(""); + await loadScope(); + } catch (e) { + setScopeErr(e instanceof Error ? e.message : "添加失败"); + } finally { + setScopeBusy(false); + } + }; + + const removeScope = async (row: TaskScopeRow) => { + setScope((prev) => prev.filter((s) => s.id !== row.id)); + try { + await api.deleteTaskScope(taskId, row.id); + } catch { + await loadScope(); // 删除失败：重新拉取还原 + } + }; const markRerun = (key: string, on: boolean) => setRerunning((prev) => { @@ -77,9 +162,7 @@ export function OverviewTab({ taskId }: { taskId: string }) { markRerun(id, true); try { await api.rerunIntent(taskId, id); - setIntents((prev) => - prev.map((i) => (i.id === id ? { ...i, state: "open" } : i)), - ); + setIntents((prev) => prev.map((i) => (i.id === id ? { ...i, state: "open" } : i))); } catch { // 失败忽略：下次轮询仍显示 blocked，用户可再点 } finally { @@ -92,9 +175,7 @@ export function OverviewTab({ taskId }: { taskId: string }) { markRerun("__all__", true); try { await api.rerunBlocked(taskId); - setIntents((prev) => - prev.map((i) => (i.state === "blocked" ? { ...i, state: "open" } : i)), - ); + setIntents((prev) => prev.map((i) => (i.state === "blocked" ? { ...i, state: "open" } : i))); } catch { // ignore } finally { @@ -107,13 +188,12 @@ export function OverviewTab({ taskId }: { taskId: string }) { const load = async () => { try { - const [tasksResp, statsResp, intentsResp, findingsResp] = - await Promise.all([ - api.tasks(), - api.stats(taskId), - api.intents(taskId), - api.findings(taskId), - ]); + const [tasksResp, statsResp, intentsResp, findingsResp] = await Promise.all([ + api.tasks(), + api.stats(taskId), + api.intents(taskId), + api.findings(taskId), + ]); if (cancelled) return; setTask(tasksResp.tasks.find((t) => t.id === taskId) ?? null); setStats(statsResp); @@ -133,20 +213,35 @@ export function OverviewTab({ taskId }: { taskId: string }) { }; load(); - const timer = setInterval(load, 3000); + void loadScope(); + void loadTokens(); + const timer = setInterval(() => { + void load(); + void loadTokens(); + }, 3000); return () => { cancelled = true; clearInterval(timer); }; - }, [taskId]); + }, [taskId, loadScope, loadTokens]); const running = intents.filter((i) => i.state === "running"); const open = intents.filter((i) => i.state === "open"); const blocked = intents.filter((i) => i.state === "blocked"); const taskFindings = findings.filter((f) => f.task_id === taskId); - const goalsPct = task?.goals_total - ? Math.round(((task.goals_met ?? 0) / task.goals_total) * 100) - : 0; + const goalsPct = task?.goals_total ? Math.round(((task.goals_met ?? 0) / task.goals_total) * 100) : 0; + // token 合计（跨全部模型），用于卡片头部总览。 + const tokenTotals = modelTokens.reduce( + (acc, m) => { + acc.input += m.input_tokens; + acc.output += m.output_tokens; + acc.cacheRead += m.cache_read_tokens; + acc.cacheWrite += m.cache_write_tokens; + acc.calls += m.calls; + return acc; + }, + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, calls: 0 }, + ); return (
@@ -160,15 +255,11 @@ export function OverviewTab({ taskId }: { taskId: string }) {
描述
-

- {task?.description?.trim() || "—"} -

+

{task?.description?.trim() || "—"}

目标
-

- {task?.goal?.trim() || "—"} -

+

{task?.goal?.trim() || "—"}

 @@ -205,6 +296,157 @@ export function OverviewTab({ taskId }: { taskId: string }) { )} + {/* LLM Token 用量：按模型分组，数据来自 llm_records（需开启 LLM 录制）。 */} + + + + LLM Token 用量 + + （按模型统计{tokenTotals.calls > 0 ? `，共 ${tokenTotals.calls} 次调用` : ""}） + + + + + {modelTokens.length > 0 ? ( + <> + {/* 合计总览 */} +
+ + 输入 + {fmtTokens(tokenTotals.input)} + + + 输出 + {fmtTokens(tokenTotals.output)} + + + 缓存读 + {fmtTokens(tokenTotals.cacheRead)} + + + 缓存命中率 + + {cacheHitRate(tokenTotals.cacheRead, tokenTotals.input)} + + +
+ {/* 按模型明细表 */} +
+
+ + + + + + + + + + + + {modelTokens.map((m) => ( + + + + + + + + + ))} + +
模型调用输入输出缓存读命中率
+ {m.model} + {m.calls}{fmtTokens(m.input_tokens)}{fmtTokens(m.output_tokens)}{fmtTokens(m.cache_read_tokens)} + {cacheHitRate(m.cache_read_tokens, m.input_tokens)} +
+
+ + ) : ( +

暂无 LLM 用量（任务尚未产生调用，或记录仍在写入）。

+ )} + + + {/* 测试范围：覆盖度分母 + 授权边界，可手动增删。 */} + + + + 测试范围 + + （覆盖度分母 + 授权边界，共 {scope.length} 条） + + + + + {/* 新增表单 */} +
+ setScopeKind(e.target.value as TaskScopeRow["kind"])} + > + 根域名 + 子域名 + IP + 网段 + 公司 + + setScopeValueInput(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") void addScope(); + }} + disabled={scopeBusy} + /> + + {scopeErr && {scopeErr}} +
+ {/* 范围列表 */} + {scope.length > 0 ? ( +
+ {scope.map((row) => ( +
+ + {SCOPE_KIND_LABELS[row.kind]} + + {scopeValue(row)} + {SCOPE_SOURCE_LABELS[row.source]} + {row.task_id.toString() === taskId ? ( + + ) : ( + 继承 + )} +
+ ))} +
+ ) : ( +

暂无测试范围，添加后可作为资产覆盖度的分母。

+ )} + + {/* Heartbeat */} @@ -224,17 +466,13 @@ export function OverviewTab({ taskId }: { taskId: string }) {
运行中 Worker
-
- {running.length} -
+
{running.length}
最近活动
- {task?.last_activity - ? new Date(task.last_activity).toLocaleTimeString("zh-CN") - : "—"} + {task?.last_activity ? new Date(task.last_activity).toLocaleTimeString("zh-CN") : "—"}
@@ -270,9 +508,7 @@ export function OverviewTab({ taskId }: { taskId: string }) { {i.payload}
))} - {running.length === 0 && ( -

暂无进行中意图

- )} + {running.length === 0 &&

暂无进行中意图

} @@ -284,27 +520,19 @@ export function OverviewTab({ taskId }: { taskId: string }) {
-
- {taskFindings.length} -
+
{taskFindings.length}
确认漏洞
-
- {running.length} -
+
{running.length}
执行中
-
- {open.length} -
+
{open.length}
frontier 待领
-
- {blocked.length} -
+
{blocked.length}
被拦意图
 @@ -323,9 +551,7 @@ export function OverviewTab({ taskId }: { taskId: string }) { {f.summary}
))} - {taskFindings.length === 0 && ( -

暂无发现

- )} + {taskFindings.length === 0 &&

暂无发现

}
@@ -337,19 +563,10 @@ export function OverviewTab({ taskId }: { taskId: string }) { 被拦/出错意图 - - （共 {blocked.length} 条，可重跑） - + （共 {blocked.length} 条，可重跑） - @@ -365,9 +582,7 @@ export function OverviewTab({ taskId }: { taskId: string }) { disabled={rerunning.has(i.id)} onClick={() => void rerunOne(i.id)} > - + 重跑
@@ -383,24 +598,9 @@ export function OverviewTab({ taskId }: { taskId: string }) { {/* Stat cards */}
- - - + + +
); diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx index 056b786..e671149 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/report-tab.tsx @@ -9,6 +9,7 @@ import { Markdown } from "@/components/markdown"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { api } from "@/lib/api"; +import { copyText } from "@/lib/utils"; export function ReportTab({ taskId }: { taskId: string }) { const [report, setReport] = React.useState(""); @@ -34,12 +35,16 @@ export function ReportTab({ taskId }: { taskId: string }) { }; }, [taskId]); - function copy() { + async function copy() { if (!report) return; - navigator.clipboard?.writeText(report); - setCopied(true); - toast.success("已复制 Markdown"); - setTimeout(() => setCopied(false), 1500); + const ok = await copyText(report); + if (ok) { + setCopied(true); + toast.success("已复制 Markdown"); + setTimeout(() => setCopied(false), 1500); + } else { + toast.error("复制失败，请手动选择文本复制"); + } } let content: React.ReactNode; diff --git a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx index 1e18d1b..b21ca81 100644 --- a/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx +++ b/web/src/app/(main)/function/tasks/detail/_tabs/sessions-tab.tsx @@ -359,8 +359,6 @@ export function SessionsTab({ taskId }: { taskId: string }) { const [olderIntents, setOlderIntents] = React.useState([]); const [intentsHasMore, setIntentsHasMore] = React.useState(false); const [loadingOlderIntents, setLoadingOlderIntents] = React.useState(false); - // Local-only chat messages overlaid onto the main-agent transcript. - const [chatExtra, setChatExtra] = React.useState([]); const [input, setInput] = React.useState(""); const [sending, setSending] = React.useState(false); const [stopping, setStopping] = React.useState(false); @@ -613,7 +611,6 @@ export function SessionsTab({ taskId }: { taskId: string }) { // ID) resumes from the DB, so a dropped realtime link self-heals; seq-merge dedups. React.useEffect(() => { setStore({}); - setChatExtra([]); setSseLive(false); setActiveId(MAIN_ID); // a stale worker id from the previous task must not leak in snapshotRef.current = 0; @@ -892,13 +889,11 @@ export function SessionsTab({ taskId }: { taskId: string }) { // worker session shows only its intent's activity, led by the intent objective. const activity = React.useMemo(() => { const items = activeState?.items ?? []; - if (isMain) { - // chatExtra is an optimistic local echo — drop any entry the persisted (SSE) - // copy already covers so messages don't show twice after the backend round-trips. - const seen = new Set(items.map((a) => `${a.kind} ${a.summary}`)); - const pending = chatExtra.filter((a) => !seen.has(`${a.kind} ${a.summary}`)); - return [...items, ...pending].sort((a, b) => a.seq - b.seq); - } + // The main-agent console renders purely from server data: the human turn is + // persisted+broadcast by the backend BEFORE it returns, so it arrives over the + // same SSE stream (worker="mainagent") as every agent step — no client-side + // optimistic echo, hence no fabricated seq that could collide with real DB ids. + if (isMain) return items; if (isPlanner || isSystem) return items; // Worker session: the intent leads the transcript as a right-aligned "user"-style // message (the task handed to this worker), followed by its execution steps. @@ -918,7 +913,6 @@ export function SessionsTab({ taskId }: { taskId: string }) { isPlanner, isSystem, activeState, - chatExtra, active.id, active.title, active.last_activity, @@ -1010,6 +1004,26 @@ export function SessionsTab({ taskId }: { taskId: string }) { const vp = viewport(); if (vp) vp.scrollTop = vp.scrollHeight; }, [activity, viewport]); + // Lazy detail loads (AnswerBlock / ToolBlock / Markdown) grow the content AFTER the + // activity array settles, WITHOUT changing its reference — so the layout effects + // above never re-fire and a freshly opened session would leave its last message + // scrolled partly off-screen (the final answer expands from a one-line summary to + // full markdown below the fold). A ResizeObserver re-pins to the bottom on any + // height change while the user is still at the bottom, so opening the main agent + // lands on the last message fully shown. contentRef's div is always mounted, so the + // observer catches the transcript mounting + each detail expanding. + // biome-ignore lint/correctness/useExhaustiveDependencies: activeId intentionally rebinds the observer to the new session's content. + React.useEffect(() => { + const el = contentRef.current; + if (!el || typeof ResizeObserver === "undefined") return; + const ro = new ResizeObserver(() => { + if (!atBottomRef.current) return; + const vp = viewport(); + if (vp) vp.scrollTop = vp.scrollHeight; + }); + ro.observe(el); + return () => ro.disconnect(); + }, [activeId, viewport]); function stop() { if (stopping) return; @@ -1021,36 +1035,19 @@ export function SessionsTab({ taskId }: { taskId: string }) { const text = input.trim(); const atts = attachments; if ((!text && atts.length === 0) || sending) return; - const base = [...(store.main?.items ?? []), ...chatExtra]; - const nextSeq = Math.max(0, ...base.map((a) => a.seq)) + 1; - const userMsg: Activity = { - seq: nextSeq, - worker: "mainagent", - ts: new Date().toISOString(), - kind: "user", - summary: text, - // 有附件时把 {text, attachments} 塞进 detail,乐观气泡直接渲染卡片,无需再拉取。 - detail: atts.length > 0 ? JSON.stringify({ text, attachments: atts }) : undefined, - }; - // optimistic echo of the human turn only; the agent's steps + final answer - // stream back live via SSE (worker="mainagent"), so we don't append the reply. - setChatExtra((prev) => [...prev, userMsg]); + // No optimistic echo: the backend persists+broadcasts the human turn before it + // returns, so it streams back over SSE (worker="mainagent") with its real DB + // seq — the transcript renders it from server data like every other step. Clear + // the composer eagerly for responsiveness; restore it if the send fails. setInput(""); setAttachments([]); setSending(true); api .chat(text, taskId, atts.length > 0 ? atts : undefined) - .catch(() => { - setChatExtra((prev) => [ - ...prev, - { - seq: nextSeq + 1, - worker: "mainagent", - ts: new Date().toISOString(), - kind: "text", - summary: "（发送失败，请稍后重试）", - }, - ]); + .catch((e) => { + setInput(text); // restore so the user doesn't lose their text / attachments + setAttachments(atts); + toast.error("发送失败：" + ((e as Error).message || "请稍后重试")); }) .finally(() => setSending(false)); } diff --git a/web/src/app/(main)/function/tasks/page.tsx b/web/src/app/(main)/function/tasks/page.tsx index ce6b5a9..bc60ba0 100644 --- a/web/src/app/(main)/function/tasks/page.tsx +++ b/web/src/app/(main)/function/tasks/page.tsx @@ -172,7 +172,11 @@ function taskDuration(task: Task, nowSec: number): number { return end > start ? end - start : 0; } -function deleteSummary(result: DeleteTaskResult): string { +// deleteDetails takes only the countable part of the result so callers can pass an +// aggregate accumulated over a bulk delete. +type DeleteCounts = Omit; + +function deleteDetails(result: DeleteCounts): string[] { const details: string[] = []; if (result.assets_deleted > 0) details.push(`删除资产 ${result.assets_deleted} 条`); if (result.assets_detached > 0) details.push(`解除共享资产关联 ${result.assets_detached} 条`); @@ -180,6 +184,11 @@ function deleteSummary(result: DeleteTaskResult): string { if (result.files_deleted) details.push("删除任务文件"); if (result.findings_deleted > 0) details.push(`删除漏洞 ${result.findings_deleted} 条`); if (result.llm_records_deleted > 0) details.push(`删除 LLM 请求/响应记录 ${result.llm_records_deleted} 条`); + return details; +} + +function deleteSummary(result: DeleteTaskResult): string { + const details = deleteDetails(result); return details.length > 0 ? `任务已删除（${details.join("，")}）` : "任务已删除"; } @@ -232,6 +241,48 @@ export default function TasksPage() { [filtered, page, pageSize], ); + // 多选删除:选择跨翻页/筛选保留,只在任务真的消失(被删或后端不再返回)时收敛。 + const [selectedIds, setSelectedIds] = React.useState>(() => new Set()); + + React.useEffect(() => { + setSelectedIds((prev) => { + if (prev.size === 0) return prev; + const live = new Set(tasks.map((t) => t.id)); + const next = new Set([...prev].filter((id) => live.has(id))); + return next.size === prev.size ? prev : next; + }); + }, [tasks]); + + const toggleSelected = React.useCallback((id: string, checked: boolean) => { + setSelectedIds((prev) => { + const next = new Set(prev); + if (checked) next.add(id); else next.delete(id); + return next; + }); + }, []); + + const pageIds = React.useMemo(() => paginated.map((t) => t.id), [paginated]); + const pageSelectedCount = React.useMemo( + () => pageIds.filter((id) => selectedIds.has(id)).length, + [pageIds, selectedIds], + ); + let headerChecked: boolean | "indeterminate" = false; + if (pageIds.length > 0 && pageSelectedCount === pageIds.length) { + headerChecked = true; + } else if (pageSelectedCount > 0) { + headerChecked = "indeterminate"; + } + + const toggleSelectedPage = React.useCallback((checked: boolean) => { + setSelectedIds((prev) => { + const next = new Set(prev); + for (const id of pageIds) { + if (checked) next.add(id); else next.delete(id); + } + return next; + }); + }, [pageIds]); + // lastRef holds the previous poll's serialized payload: the list is re-fetched every // POLL_MS but usually comes back unchanged, and setTasks on an identical payload would // re-render the whole page for nothing. Bail out when it matches. @@ -284,6 +335,63 @@ export default function TasksPage() { } }, [load]); + // deleteTasks 逐个删除所选任务:后端没有批量接口,且单次删除会连带清理资产/流量/文件, + // 串行执行以免一次性打爆后端;成功的从选中集移除,失败的保留以便重试。 + const deleteTasks = React.useCallback(async ( + ids: string[], + options: DeleteTaskOptions, + onProgress: (done: number) => void, + ) => { + const total: DeleteCounts = { + assets_deleted: 0, + assets_detached: 0, + traffic_deleted: 0, + files_deleted: false, + findings_deleted: 0, + llm_records_deleted: 0, + }; + const deleted: string[] = []; + const failed: { id: string; message: string }[] = []; + const warnings: string[] = []; + + for (const id of ids) { + try { + const r = await api.deleteTask(id, options); + total.assets_deleted += r.assets_deleted; + total.assets_detached += r.assets_detached; + total.traffic_deleted += r.traffic_deleted; + total.files_deleted = total.files_deleted || r.files_deleted; + total.findings_deleted += r.findings_deleted; + total.llm_records_deleted += r.llm_records_deleted; + if (r.cleanup_warning) warnings.push(`#${id}：${r.cleanup_warning}`); + deleted.push(id); + } catch (e) { + failed.push({ id, message: (e as Error).message }); + } + onProgress(deleted.length + failed.length); + } + + if (deleted.length > 0) { + setSelectedIds((prev) => { + const next = new Set(prev); + for (const id of deleted) next.delete(id); + return next; + }); + const details = deleteDetails(total); + const summary = `已删除 ${deleted.length} 个任务` + (details.length > 0 ? `（${details.join("，")}）` : ""); + if (warnings.length > 0) { + toast.warning(`${summary}；部分外部数据清理未完成：${warnings.join("；")}`); + } else { + toast.success(summary); + } + } + if (failed.length > 0) { + const head = failed.slice(0, 3).map((f) => `#${f.id}（${f.message}）`).join("；"); + toast.error(`${failed.length} 个任务删除失败：${head}${failed.length > 3 ? " 等" : ""}`); + } + load(); + }, [load]); + return ( @@ -325,6 +433,15 @@ export default function TasksPage() { {filtered.length}/{tasks.length} 条 + {selectedIds.size > 0 && ( + <> + 已选 {selectedIds.size} 个 + + + + )}
@@ -341,6 +458,13 @@ export default function TasksPage() { + + toggleSelectedPage(checked === true)} + aria-label="选择本页全部任务" + /> + ID 描述 目标 @@ -361,6 +485,8 @@ export default function TasksPage() { // 带来的整表重渲染,只让在跑的那几行走时长。 nowSec={task.status === "running" ? nowSec : 0} onDelete={deleteTask} + selected={selectedIds.has(task.id)} + onSelectedChange={toggleSelected} /> ))} @@ -461,13 +587,24 @@ const TaskRow = React.memo(function TaskRow({ task, nowSec, onDelete, + selected, + onSelectedChange, }: { task: Task; nowSec: number; onDelete: (id: string, options: DeleteTaskOptions) => Promise; + selected: boolean; + onSelectedChange: (id: string, checked: boolean) => void; }) { return ( - + + + onSelectedChange(task.id, checked === true)} + aria-label={`选择任务 ${task.id}`} + /> + {task.id} @@ -554,6 +691,117 @@ const deleteOptionKeys: (keyof DeleteTaskOptions)[] = [ "delete_llm_records", ]; +// DeleteOptionFields renders the「同时清理关联数据」checkbox block shared by the +// single-task and bulk delete dialogs. idPrefix keeps the label/input ids unique when +// several dialogs live in the same table. +function DeleteOptionFields({ + idPrefix, + options, + onOptionsChange, + disabled, +}: { + idPrefix: string; + options: DeleteTaskOptions; + onOptionsChange: React.Dispatch>; + disabled: boolean; +}) { + const selectedCount = deleteOptionKeys.filter((key) => options[key]).length; + let allChecked: boolean | "indeterminate" = false; + if (selectedCount === deleteOptionKeys.length) { + allChecked = true; + } else if (selectedCount > 0) { + allChecked = "indeterminate"; + } + + const updateOption = (key: keyof DeleteTaskOptions, checked: boolean) => { + onOptionsChange((current) => ({ ...current, [key]: checked })); + }; + + const updateAllOptions = (checked: boolean) => { + onOptionsChange({ + delete_assets: checked, + delete_traffic: checked, + delete_files: checked, + delete_findings: checked, + delete_llm_records: checked, + }); + }; + + return ( +
+ 同时清理关联数据 + + + updateAllOptions(checked === true)} + /> + + 全部删除 + 选中下方全部关联数据，包括资产、流量、文件、漏洞和 LLM 请求/响应记录。 + + + + updateOption("delete_assets", checked === true)} + /> + + 关联资产 + 删除仅属于该任务的资产；共享资产只解除当前任务关联。 + + + + updateOption("delete_traffic", checked === true)} + /> + + 关联流量 + 按关联资产的精确主机名删除；仍被其他任务引用的共享主机流量会保留。 + + + + updateOption("delete_files", checked === true)} + /> + + 任务文件 + 删除该任务工作目录中的上传文件、命令输出和其他产物。 + + + + updateOption("delete_findings", checked === true)} + /> + + 关联漏洞 + 永久删除该任务产生的独立漏洞记录与漏洞报告。 + + + + updateOption("delete_llm_records", checked === true)} + /> + + LLM 请求/响应记录 + 永久删除该任务录制的 LLM 请求、响应、Token 与错误详情。 + + + +
+ ); +} + function DeleteTaskDialog({ task, onDelete, @@ -565,29 +813,6 @@ function DeleteTaskDialog({ const [deleting, setDeleting] = React.useState(false); const [options, setOptions] = React.useState(emptyDeleteOptions); - const selectedCount = deleteOptionKeys.filter((key) => options[key]).length; - const allSelected = selectedCount === deleteOptionKeys.length; - let allChecked: boolean | "indeterminate" = false; - if (allSelected) { - allChecked = true; - } else if (selectedCount > 0) { - allChecked = "indeterminate"; - } - - const updateOption = (key: keyof DeleteTaskOptions, checked: boolean) => { - setOptions((current) => ({ ...current, [key]: checked })); - }; - - const updateAllOptions = (checked: boolean) => { - setOptions({ - delete_assets: checked, - delete_traffic: checked, - delete_files: checked, - delete_findings: checked, - delete_llm_records: checked, - }); - }; - const handleOpenChange = (next: boolean) => { if (deleting) return; setOpen(next); @@ -621,77 +846,12 @@ function DeleteTaskDialog({ 的执行记录与探索链路将被永久删除。 -
- 同时清理关联数据 - - - updateAllOptions(checked === true)} - /> - - 全部删除 - 选中下方全部关联数据，包括资产、流量、文件、漏洞和 LLM 请求/响应记录。 - - - - updateOption("delete_assets", checked === true)} - /> - - 关联资产 - 删除仅属于该任务的资产；共享资产只解除当前任务关联。 - - - - updateOption("delete_traffic", checked === true)} - /> - - 关联流量 - 按关联资产的精确主机名删除；仍被其他任务引用的共享主机流量会保留。 - - - - updateOption("delete_files", checked === true)} - /> - - 任务文件 - 删除该任务工作目录中的上传文件、命令输出和其他产物。 - - - - updateOption("delete_findings", checked === true)} - /> - - 关联漏洞 - 永久删除该任务产生的独立漏洞记录与漏洞报告。 - - - - updateOption("delete_llm_records", checked === true)} - /> - - LLM 请求/响应记录 - 永久删除该任务录制的 LLM 请求、响应、Token 与错误详情。 - - - -
+ 取消 void, + ) => Promise; +}) { + const [open, setOpen] = React.useState(false); + const [deleting, setDeleting] = React.useState(false); + const [done, setDone] = React.useState(0); + const [options, setOptions] = React.useState(emptyDeleteOptions); + + const handleOpenChange = (next: boolean) => { + if (deleting) return; + setOpen(next); + if (next) { + setOptions(emptyDeleteOptions()); + setDone(0); + } + }; + + const handleDelete = async () => { + setDeleting(true); + setDone(0); + try { + await onDelete(ids, options, setDone); + setOpen(false); + } finally { + setDeleting(false); + } + }; + + return ( + + + + + + + 确认删除所选 {ids.length} 个任务？ + + 这些任务的执行记录与探索链路将被永久删除，下方清理选项对所选任务统一生效。 + + +
+ {ids.slice(0, 30).map((id) => ( + #{id} + ))} + {ids.length > 30 && …等 {ids.length} 个} +
+ + + 取消 + { + event.preventDefault(); + void handleDelete(); + }} + > + {deleting && } + {deleting ? `删除中 ${done}/${ids.length}` : `删除 ${ids.length} 个任务`} + + + + + ); +} + // CreateTaskSheet is the 新建任务 drawer. Its form state lives HERE, not in TasksPage: with // 描述/目标 held by the page component every keystroke re-rendered the whole task table // behind the drawer (plus its sticky column and 20 AlertDialog trees), which showed up as @@ -872,7 +1116,7 @@ function CreateTaskSheet({ tasks, onCreated }: { tasks: Task[]; onCreated: () => return ( - diff --git a/web/src/app/(main)/system/llm/page.tsx b/web/src/app/(main)/system/llm/page.tsx index fb37607..5ac1e28 100644 --- a/web/src/app/(main)/system/llm/page.tsx +++ b/web/src/app/(main)/system/llm/page.tsx @@ -29,26 +29,26 @@ import { api } from "@/lib/api"; import type { LLMPoolMember, LLMPoolStatus, LLMProfile } from "@/lib/types"; import { cn } from "@/lib/utils"; -// 思考强度档位（仅在开关打开时生效）。 +// 思考开关(thinking.type)与思考强度(reasoning_effort)是两个【互相独立】的字段， +// 各自单独设置——有些接口没有 thinking 字段、只靠强度参数就能激活思考，故需解耦。 +// 存库空字符串 = 该字段【不发送】；Radix Select 不接受空 value，故 UI 用 "none" +// 哨兵表示不发送，存取时与 "" 互转（NONE / fromStore / toStore）。 +const NONE = "none"; +const fromStore = (v?: string) => (v ? v : NONE); +const toStore = (v: string) => (v === NONE ? "" : v); +const THINKING_TYPES: { value: string; label: string }[] = [ + { value: NONE, label: "不发送（默认）" }, + { value: "disabled", label: "关闭" }, + { value: "enabled", label: "开启" }, +]; const EFFORT_LEVELS: { value: string; label: string }[] = [ - { value: "low", label: "低" }, - { value: "medium", label: "中" }, - { value: "high", label: "高" }, - { value: "max", label: "最大" }, + { value: NONE, label: "不发送（默认）" }, + { value: "low", label: "low" }, + { value: "medium", label: "medium" }, + { value: "high", label: "high" }, + { value: "xhigh", label: "xhigh" }, + { value: "max", label: "max" }, ]; -// 思考模式三态 ↔ 存库字符串 reasoning_effort： -// none→""（thinking / reasoning_effort 两个字段都【不发送】，兼容不支持该字段的模型，如 MiniMax） -// off →"off"（发送 thinking:{type:disabled}，用于默认开启思考、需要显式关闭的模型） -// on →强度值（发送 thinking:{enabled} + reasoning_effort） -type ThinkMode = "none" | "off" | "on"; -const THINK_MODES: { value: ThinkMode; label: string }[] = [ - { value: "none", label: "不发送（默认）" }, - { value: "off", label: "显式关闭" }, - { value: "on", label: "开启" }, -]; -const toStore = (mode: ThinkMode, effort: string) => (mode === "on" ? effort : mode === "off" ? "off" : ""); -const modeFromStore = (v?: string): ThinkMode => (!v ? "none" : v === "off" ? "off" : "on"); -const effortFromStore = (v?: string) => (v && v !== "off" ? v : "high"); function cooldownText(secs: number) { if (secs <= 0) return ""; @@ -305,8 +305,8 @@ function ProfileSheet({ const [rps, setRps] = React.useState("0"); const [rpm, setRpm] = React.useState("0"); const [cw, setCw] = React.useState("0"); // 上下文窗口(K tokens);0=默认200K - const [thinkMode, setThinkMode] = React.useState("none"); - const [effort, setEffort] = React.useState("high"); + const [thinkingType, setThinkingType] = React.useState(NONE); + const [effort, setEffort] = React.useState(NONE); const [priority, setPriority] = React.useState("0"); // 轮询顺位;越大越先 const [poolExclude, setPoolExclude] = React.useState(false); const [testing, setTesting] = React.useState(false); @@ -327,8 +327,8 @@ function ProfileSheet({ setRps(String(profile?.rate_per_second ?? 0)); setRpm(String(profile?.rate_per_minute ?? 0)); setCw(String(profile?.context_window_k ?? 0)); - setThinkMode(modeFromStore(profile?.reasoning_effort)); - setEffort(effortFromStore(profile?.reasoning_effort)); + setThinkingType(fromStore(profile?.thinking_type)); + setEffort(fromStore(profile?.reasoning_effort)); setPriority(String(profile?.priority ?? 0)); setPoolExclude(profile?.pool_exclude ?? false); setApiKey(""); @@ -365,7 +365,7 @@ function ProfileSheet({ try { // 用配置实际会跑的思考参数来测，这样不支持该字段的模型在这里就失败， // 而不是等到跑任务时才炸。传 profile id：Key 输入框留空时用已存的 Key。 - const r = await api.testLLM(format, model, baseUrl, apiKey, proxy, toStore(thinkMode, effort), profileId); + const r = await api.testLLM(format, model, baseUrl, apiKey, proxy, toStore(thinkingType), toStore(effort), profileId); if (r.ok) toast.success(`连接成功 · ${r.latency_ms ?? "?"}ms · ${r.model ?? model}`); else toast.error(`连接失败：${r.error ?? "未知"}`); } catch (e) { @@ -394,7 +394,8 @@ function ProfileSheet({ rate_per_second: Number(rps) || 0, rate_per_minute: Number(rpm) || 0, context_window_k: Number(cw) || 0, - reasoning_effort: toStore(thinkMode, effort), + thinking_type: toStore(thinkingType), + reasoning_effort: toStore(effort), priority: Number(priority) || 0, pool_exclude: poolExclude, }); @@ -601,18 +602,18 @@ function ProfileSheet({
- +

- 不发送=不带思考字段（兼容 MiniMax 等不支持该字段的模型）；显式关闭=发 - disabled（默认开思考的模型才需要）；开启=先推理再作答。 + 控制是否发送 thinking 字段。不发送=不带该字段（兼容 MiniMax 等不支持 + 的模型）；关闭=发 disabled；开启=发 enabled。与下面的强度互相独立。

- - {THINK_MODES.map((o) => ( + {THINKING_TYPES.map((o) => ( {o.label} @@ -620,28 +621,27 @@ function ProfileSheet({
- {thinkMode === "on" && ( -
-
- -

- 推理投入档位，越高越深入（OpenAI reasoning_effort / Anthropic output_config.effort）。 -

-
- +
+
+ +

+ 独立的强度档位（OpenAI reasoning_effort / Anthropic output_config.effort）。 + 有些接口没有 thinking 字段、只靠强度即可激活思考，故可单独设置、不发送思考开关。 +

- )} + +
diff --git a/web/src/app/(main)/system/skills/page.tsx b/web/src/app/(main)/system/skills/page.tsx index ab6366d..41539df 100644 --- a/web/src/app/(main)/system/skills/page.tsx +++ b/web/src/app/(main)/system/skills/page.tsx @@ -12,13 +12,25 @@ import { PlusIcon, Trash2Icon, UploadIcon, + AlertTriangleIcon, } from "lucide-react"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Textarea } from "@/components/ui/textarea"; import { Checkbox } from "@/components/ui/checkbox"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Badge } from "@/components/ui/badge"; import { @@ -31,7 +43,17 @@ import { import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { cn } from "@/lib/utils"; import { api } from "@/lib/api"; -import type { Agent, SkillItem, MCPServer } from "@/lib/types"; +import type { Agent, SkillItem, MCPServer, SkillCall, MissingSkill } from "@/lib/types"; + +function fmtTime(ts?: string) { + if (!ts) return "从未调用"; + return new Date(ts).toLocaleString("zh-CN", { + month: "2-digit", + day: "2-digit", + hour: "2-digit", + minute: "2-digit", + }); +} // ── Tree node ────────────────────────────────────────────────────────────── interface TreeNode { @@ -85,9 +107,22 @@ function buildTree(entries: string[]): TreeNode[] { if (fname) nodes.push({ name: fname, path: entry, type: "file", children: [] }); } } + sortNodes(root); return root; } +// sortNodes orders every level like a file explorer: directories before files, +// then case-insensitive by name. Recurses into children. +function sortNodes(nodes: TreeNode[]): void { + nodes.sort((a, b) => { + if (a.type !== b.type) return a.type === "dir" ? -1 : 1; + return a.name.localeCompare(b.name, undefined, { sensitivity: "base" }); + }); + for (const n of nodes) { + if (n.children.length > 0) sortNodes(n.children); + } +} + // ── State types ─────────────────────────────────────────────────────────── type Selected = | { skill: string; path: null } @@ -99,6 +134,164 @@ type Creating = { kind: "file" | "dir"; } | null; +type PendingDelete = + | { kind: "skill"; skill: string } + | { kind: "file"; skill: string; path: string } + | { kind: "dir"; skill: string; path: string } + | null; + +// ── Overview (empty-state) ────────────────────────────────────────────────── +// Shown when nothing is selected: a library-wide snapshot from data already loaded +// (the skill list carries per-skill usage; missing is fetched alongside). No extra +// requests — this is pure aggregation over props. +function SkillsOverview({ + skills, + missing, + onSelect, +}: { + skills: SkillItem[]; + missing: MissingSkill[]; + onSelect: (name: string) => void; +}) { + const agg = React.useMemo(() => { + const totalCalls = skills.reduce((n, s) => n + s.calls, 0); + const used = skills.filter((s) => s.calls > 0); + const ranked = [...used].sort((a, b) => b.calls - a.calls); + const neverUsed = skills.filter((s) => s.calls === 0); + const recent = skills + .filter((s) => s.last_used) + .sort((a, b) => (a.last_used! < b.last_used! ? 1 : -1)) + .slice(0, 6); + const missingCalls = missing.reduce((n, m) => n + m.calls, 0); + return { + totalCalls, + usedCount: used.length, + ranked, + topCalls: ranked[0]?.calls ?? 0, + neverUsed, + recent, + missingCalls, + }; + }, [skills, missing]); + + if (skills.length === 0) { + return ( +
+

暂无 Skill，点击左侧「新建」或「上传压缩包」开始

+
+ ); + } + + const stats: { label: string; value: React.ReactNode; hint?: string }[] = [ + { label: "Skill 总数", value: skills.length, hint: `${agg.usedCount} 个被调用过` }, + { label: "累计调用", value: agg.totalCalls }, + { label: "未使用", value: agg.neverUsed.length, hint: agg.neverUsed.length > 0 ? "从未被任何 agent 加载" : "全部用过" }, + { label: "未命中调用", value: agg.missingCalls, hint: missing.length > 0 ? `${missing.length} 个不存在的 skill` : "无" }, + ]; + + return ( +
+
+

技能库总览

+

选择左侧的 Skill 查看详情与调用记录，或从这里快速了解整体使用情况。

+
+ + {/* 指标卡 */} +
+ {stats.map((s) => ( +
+

{s.value}

+

{s.label}

+ {s.hint &&

{s.hint}

} +
+ ))} +
+ + {/* 调用排行 */} +
+ + {agg.ranked.length === 0 ? ( +

+ 还没有任何 Skill 调用记录。 +

+ ) : ( +
+ {agg.ranked.slice(0, 8).map((s) => ( + + ))} +
+ )} +
+ +
+ {/* 最近调用 */} +
+ + {agg.recent.length === 0 ? ( +

暂无记录。

+ ) : ( +
+ {agg.recent.map((s) => ( + + ))} +
+ )} +
+ + {/* 未使用（可清理 / 需曝光） */} +
+ + {agg.neverUsed.length === 0 ? ( +

所有 Skill 都被调用过。

+ ) : ( +
+ {agg.neverUsed.map((s) => ( + + ))} +
+ )} +
+
+
+ ); +} + // ── Page ────────────────────────────────────────────────────────────────── export default function SkillsPage() { const [skills, setSkills] = React.useState([]); @@ -137,10 +330,21 @@ export default function SkillsPage() { // skill detail — MCP edit state (optimistic, rolls back on error) const [detailMcps, setDetailMcps] = React.useState([]); + // delete confirmation + const [pendingDelete, setPendingDelete] = React.useState(null); + const [deleting, setDeleting] = React.useState(false); + + // 调用统计：列表页的次数/最近调用随 api.skills() 一起回来；选中某个 skill 时再拉它的 + // 最近调用明细。missing = 被点名但不存在的 skill（想用但没有）。 + const [usageCalls, setUsageCalls] = React.useState([]); + const [usageLoading, setUsageLoading] = React.useState(false); + const [missing, setMissing] = React.useState([]); + // ── Data ────────────────────────────────────────────────────────────────── const load = React.useCallback(() => { api.agents().then(setAgents).catch(() => {}); api.mcpServers().then(setMcpOptions).catch(() => {}); + api.missingSkills().then(setMissing).catch(() => {}); api.skills().then((ss) => { setSkills(ss); ss.forEach((s) => @@ -195,6 +399,17 @@ export default function SkillsPage() { setDetailMcps(sk?.mcps ?? []); }, [selected, skills]); + // Recent calls for the selected skill (detail panel only). + React.useEffect(() => { + if (!selected || selected.path !== null) { setUsageCalls([]); return; } + const name = selected.skill; + setUsageLoading(true); + api.skillUsage(name, 20) + .then((calls) => setUsageCalls(calls)) + .catch(() => setUsageCalls([])) + .finally(() => setUsageLoading(false)); + }, [selected]); + React.useEffect(() => { if (!selected || selected.path === null) { setFileContent(""); setDirty(false); return; } setFileLoading(true); @@ -296,6 +511,19 @@ export default function SkillsPage() { } } + async function runPendingDelete() { + const p = pendingDelete; + if (!p) return; + setDeleting(true); + try { + if (p.kind === "skill") await deleteSkill(p.skill); + else await deletePath(p.skill, p.path); + } finally { + setDeleting(false); + setPendingDelete(null); + } + } + async function saveFile() { if (!selected || selected.path === null) return; setSaving(true); @@ -392,7 +620,9 @@ export default function SkillsPage() { // ── Recursive tree renderer ─────────────────────────────────────────────── function renderTree(nodes: TreeNode[], skill: string, depth: number): React.ReactNode { - const baseIndent = 8 + depth * 14; + // depth+1: the skill root sits at 8px (px-2); its children indent one step + // further in so the tree reads as nested under the skill folder. + const baseIndent = 8 + (depth + 1) * 14; return nodes.map((node) => { if (node.type === "dir") { const key = `${skill}:${node.path}`; @@ -400,7 +630,7 @@ export default function SkillsPage() { return (
toggleExpanded(key)} > @@ -409,8 +639,9 @@ export default function SkillsPage() { ? : } - {node.name} - + {node.name} + {/* Absolute so a long name can never push the actions out of view */} + @@ -442,19 +673,24 @@ export default function SkillsPage() {
setSelected({ skill, path: node.path })} > - {node.name} - + {node.name} +
); }); @@ -466,9 +702,36 @@ export default function SkillsPage() { return (
-
-

Skill

-

技能库 · agentskills.io 规范 · 按 Agent 授权可见

+
+
+

Skill

+

技能库 · agentskills.io 规范 · 按 Agent 授权可见

+
+ {/* 缺口清单：agent 点名调用、但库里没有的 skill —— 直接是该补什么的依据。 */} + {missing.length > 0 && ( + + + + + +

+ Agent 点名调用、但技能库里不存在的 skill。按被点名次数排序。 +

+
+ {missing.map((m) => ( +
+ {m.skill} + {m.calls} 次 + {fmtTime(m.last_used)} +
+ ))} +
+ + + )}
{/* ── 左侧文件树 ── */} @@ -496,7 +759,10 @@ export default function SkillsPage() { {uploading ? "上传中…" : "上传压缩包"}
- + {/* Radix viewport wraps children in a display:table div that grows with + content — force it to block so long names truncate instead of + widening the rows past the sidebar. */} +
{skills.map((s) => { const isOpen = expanded.has(s.name); @@ -507,7 +773,7 @@ export default function SkillsPage() { {/* skill 根节点 */}
{ @@ -520,8 +786,19 @@ export default function SkillsPage() { ? : } - {s.name} - + {s.name} + {s.calls > 0 && ( + + {s.calls} + + )} + @@ -559,13 +836,15 @@ export default function SkillsPage() { {/* ── 右侧面板 ── */}
{!selected && ( -
-

选择左侧 Skill 或文件

-
+ setSelected({ skill: name, path: null })} + /> )} {selected && selected.path === null && selectedSkill && ( -
+

{selectedSkill.name}

{selectedSkill.description && ( @@ -581,45 +860,100 @@ export default function SkillsPage() {
- {/* ── 关联 MCP ── */} -
- - {mcpOptions.length === 0 ? ( -

暂无 MCP，可在「MCP」页添加。

- ) : ( -
- {mcpOptions.map((m) => ( - - ))} + {/* 左右分栏：配置（MCP/可见性）在左为主，调用统计在右为辅。 + lg 以下放不下时用 flex-row-reverse 回落到单列——统计因 DOM 顺序在前， + 窄屏时自然落到配置上方（与改版前的上下顺序一致）。 */} +
+ {/* ── 右侧：调用统计 ── */} +
+ +
+
+

{selectedSkill.calls}

+

总调用次数

+
+
+

{selectedSkill.tasks}

+

覆盖任务数

+
+
+

+ {fmtTime(selectedSkill.last_used)} +

+

最近调用

+
- )} -
- - {/* ── 可见性 ── */} -
- -
- {agents.map((a) => ( - - ))} - {agents.length === 0 && ( - （暂无 Agent） + {selectedSkill.usage_agents.length > 0 && ( +
+ 调用方： + {selectedSkill.usage_agents.map((k) => ( + {k} + ))} +
)} + {usageLoading ? ( +

加载调用明细…

+ ) : usageCalls.length > 0 ? ( +
+
最近 {usageCalls.length} 次调用
+
+ {usageCalls.map((c, i) => ( +
+ {fmtTime(c.ts)} + {c.agent_key || "—"} + + {c.task_id > 0 ? `任务 #${c.task_id}` : c.session_id ? "对话会话" : "—"} + +
+ ))} +
+
+ ) : ( +

还没有调用记录。

+ )} +
+ + {/* ── 左侧：关联 MCP + 可见性 ── */} +
+
+ + {mcpOptions.length === 0 ? ( +

暂无 MCP，可在「MCP」页添加。

+ ) : ( +
+ {mcpOptions.map((m) => ( + + ))} +
+ )} +
+ +
+ +
+ {agents.map((a) => ( + + ))} + {agents.length === 0 && ( + （暂无 Agent） + )} +
+
@@ -649,6 +983,35 @@ export default function SkillsPage() {
+ {/* ── 删除二次确认 ── */} + { if (!o) setPendingDelete(null); }}> + + + + {pendingDelete?.kind === "skill" && `删除 Skill「${pendingDelete.skill}」？`} + {pendingDelete?.kind === "dir" && `删除文件夹「${pendingDelete.path}」？`} + {pendingDelete?.kind === "file" && `删除文件「${pendingDelete.path}」？`} + + + {pendingDelete?.kind === "skill" + ? "将删除该 Skill 的全部文件、MCP 关联与可见性配置。此操作不可撤销。" + : pendingDelete?.kind === "dir" + ? "将一并删除该文件夹下的所有文件。此操作不可撤销。" + : "此操作不可撤销。"} + + + + 取消 + { e.preventDefault(); void runPendingDelete(); }} + > + {deleting ? "删除中…" : "删除"} + + + + + {/* ── 新建 Skill 对话框 ── */} diff --git a/web/src/components/copy-button.tsx b/web/src/components/copy-button.tsx new file mode 100644 index 0000000..6949a9e --- /dev/null +++ b/web/src/components/copy-button.tsx @@ -0,0 +1,67 @@ +"use client"; + +import * as React from "react"; + +import { CheckIcon, CopyIcon } from "lucide-react"; +import { toast } from "sonner"; + +import { Button } from "@/components/ui/button"; +import { cn, copyText } from "@/lib/utils"; + +type CopyButtonProps = { + // 要复制的文本;为空则按钮禁用。 + text: string | null | undefined; + // 复制成功后的 toast 文案,默认「已复制」。 + successMessage?: string; + label?: React.ReactNode; + size?: React.ComponentProps["size"]; + variant?: React.ComponentProps["variant"]; + className?: string; +}; + +// CopyButton 统一的「复制到剪贴板」按钮:内置成功/失败反馈,并在 HTTP 非安全上下文 +// 下自动降级(见 copyText)。 +export function CopyButton({ + text, + successMessage = "已复制", + label = "复制", + size = "sm", + variant = "outline", + className, +}: CopyButtonProps) { + const [copied, setCopied] = React.useState(false); + const timer = React.useRef | null>(null); + + React.useEffect(() => { + return () => { + if (timer.current) clearTimeout(timer.current); + }; + }, []); + + async function handleCopy() { + if (!text) return; + const ok = await copyText(text); + if (ok) { + setCopied(true); + toast.success(successMessage); + if (timer.current) clearTimeout(timer.current); + timer.current = setTimeout(() => setCopied(false), 1500); + } else { + toast.error("复制失败，请手动选择文本复制"); + } + } + + return ( + + ); +} diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index 9b8e611..6e4bc40 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -6,15 +6,63 @@ import { MOCK } from "@/lib/mock/enabled"; import { mockHandle } from "@/lib/mock/handler"; import type { - Task, DeleteTaskOptions, DeleteTaskResult, Stats, Asset, AssetNode, Edge, Company, TaskNode, Finding, Severity, FindingStatus, - FindingsPage, FindingStats, FindingQuery, Activity, - Audit, TrafficResp, TrafficDetail, TrafficHost, Settings, LLMProfile, Agent, AgentDetail, PromptVar, - PromptVersion, MCPServer, MCPTool, SkillItem, TokenUsage, TokenTotal, DailyTokenBucket, - Tool, Conversation, AgentTrigger, ChatAttachment, - InterceptRule, InterceptPending, InterceptApprovalRow, - TaskAssetView, SSProject, SSTask, ConvTokenSummary, CoverageGraphData, CoverageAssetRefs, - WorkspaceListing, WorkspaceFile, - CommandRecord, LLMRecordItem, LLMRecordDetail, LLMTask, LLMPoolStatus, + Activity, + Agent, + AgentDetail, + AgentTrigger, + Asset, + AssetNode, + Audit, + ChatAttachment, + CommandRecord, + Company, + Conversation, + ConvTokenSummary, + CoverageAssetRefs, + CoverageGraphData, + DailyTokenBucket, + DeleteTaskOptions, + DeleteTaskResult, + Edge, + Finding, + FindingQuery, + FindingStats, + FindingStatus, + FindingsPage, + InterceptApprovalRow, + InterceptPending, + InterceptRule, + LLMPoolStatus, + LLMProfile, + LLMRecordDetail, + LLMRecordItem, + LLMTask, + MCPServer, + MCPTool, + ModelTokenStat, + PromptVar, + PromptVersion, + Settings, + Severity, + SkillCall, + SkillItem, + MissingSkill, + SSProject, + SSTask, + Stats, + Task, + TaskAssetView, + TaskNode, + TaskScopeRow, + TokenTotal, + TokenUsage, + Tool, + TrafficDetail, + TrafficHost, + TrafficResp, + UsageStats, + WorkspaceFile, + WorkspaceListing, } from "@/lib/types"; function getToken(): string | null { @@ -92,19 +140,21 @@ function mockReport(_task?: string): string { export function sseUrl(path: string): string { const base = process.env.NEXT_PUBLIC_SSE_BASE ?? - (typeof window !== "undefined" - ? `${window.location.protocol}//${window.location.hostname}:8787` - : ""); + (typeof window !== "undefined" ? `${window.location.protocol}//${window.location.hostname}:8787` : ""); const token = getToken(); const sep = path.includes("?") ? "&" : "?"; return token ? `${base}${path}${sep}token=${encodeURIComponent(token)}` : `${base}${path}`; } const get = (p: string) => http(p); -const post = (p: string, body?: unknown) => http(p, { method: "POST", body: body ? JSON.stringify(body) : undefined }); -const put = (p: string, body?: unknown) => http(p, { method: "PUT", body: body ? JSON.stringify(body) : undefined }); -const patch = (p: string, body?: unknown) => http(p, { method: "PATCH", body: body ? JSON.stringify(body) : undefined }); -const del = (p: string, body?: unknown) => http(p, { method: "DELETE", body: body ? JSON.stringify(body) : undefined }); +const post = (p: string, body?: unknown) => + http(p, { method: "POST", body: body ? JSON.stringify(body) : undefined }); +const put = (p: string, body?: unknown) => + http(p, { method: "PUT", body: body ? JSON.stringify(body) : undefined }); +const patch = (p: string, body?: unknown) => + http(p, { method: "PATCH", body: body ? JSON.stringify(body) : undefined }); +const del = (p: string, body?: unknown) => + http(p, { method: "DELETE", body: body ? JSON.stringify(body) : undefined }); // Go serializes nil slices as JSON null — coerce to []. const arr = (x: T[] | null | undefined): T[] => x ?? []; @@ -116,15 +166,14 @@ export const api = { // ---- auth ---- authStatus: () => get<{ initialized: boolean }>("/auth/status"), - login: (username: string, password: string) => - post<{ token: string }>("/auth/login", { username, password }), - initPassword: (password: string) => - post<{ token: string }>("/auth/init", { password }), + login: (username: string, password: string) => post<{ token: string }>("/auth/login", { username, password }), + initPassword: (password: string) => post<{ token: string }>("/auth/init", { password }), changePassword: (oldPassword: string, newPassword: string) => post<{ ok: boolean }>("/auth/change-password", { old_password: oldPassword, new_password: newPassword }), // ---- tasks ---- - tasks: () => get<{ tasks: Task[]; active: string }>("/tasks").then((r) => ({ tasks: arr(r.tasks), active: r.active ?? "" })), + tasks: () => + get<{ tasks: Task[]; active: string }>("/tasks").then((r) => ({ tasks: arr(r.tasks), active: r.active ?? "" })), createTask: (input: { description: string; goal: string; @@ -155,9 +204,9 @@ export const api = { llm_profile_ids: llmProfileIds, active_llm_profile_id: activeLLMProfileId ?? null, }), - deleteTask: (id: string, options: DeleteTaskOptions) => - del(`/tasks/${id}`, options), - controlTask: (id: string, action: "pause" | "resume") => post<{ id: string; paused: boolean }>(`/tasks/${id}/control`, { action }), + deleteTask: (id: string, options: DeleteTaskOptions) => del(`/tasks/${id}`, options), + controlTask: (id: string, action: "pause" | "resume") => + post<{ id: string; paused: boolean }>(`/tasks/${id}/control`, { action }), controlIntent: (taskId: string, intentId: string, action: "pause" | "resume" | "cancel") => post<{ id: number; @@ -165,7 +214,8 @@ export const api = { deleted?: { intents: number; facts: number; findings: number; activities: number }; }>(`/tasks/${taskId}/intents/${intentId}/control`, { action }), // 重跑一条没跑成功的意图(blocked/exhausted/stopped)：置回 open，worker 会重新认领、从头再跑。 - rerunIntent: (taskId: string, intentId: string) => post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/${intentId}/rerun`), + rerunIntent: (taskId: string, intentId: string) => + post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/${intentId}/rerun`), // 批量重跑本任务全部 blocked 意图（一次网络/LLM 断连导致多条 blocked 时一键全部重试）。 rerunBlocked: (taskId: string) => post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/rerun-blocked`), setActive: (id: string) => post<{ active: string }>("/active", { id }), @@ -181,23 +231,26 @@ export const api = { by_type: { type: string; total: number; tested: number }[]; }>(`/tasks/${id}/coverage`), // 资产覆盖图：范围内全部资产 + 连接用的根域名/公司节点，含 tested/in_scope。 - taskCoverageGraph: (id: string) => - get(`/tasks/${id}/coverage-graph`), + taskCoverageGraph: (id: string) => get(`/tasks/${id}/coverage-graph`), + // 全局 llm_usage 聚合（仪表盘新版 token 视图）：按 profile 总量 + 按天分桶。 + usageStats: (days = 365) => get(`/tokens/usage?days=${days}`), + // 本任务测试范围列表（含继承自来源任务的范围）。 + taskScope: (id: string) => get<{ scope: TaskScopeRow[] }>(`/tasks/${id}/scope`), + // 手动新增一条测试范围（kind=company/root_domain/subdomain/ip/cidr）。 + addTaskScope: (id: string, kind: string, value: string, reason?: string) => + post(`/tasks/${id}/scope`, { kind, value, reason }), + // 删除本任务的一条测试范围。 + deleteTaskScope: (id: string, scopeId: number) => del<{ ok: boolean }>(`/tasks/${id}/scope/${scopeId}`), // 某资产在本任务里关联的意图 / 事实 / 发现（覆盖图节点抽屉用）。 - taskAssetRefs: (id: string, assetId: number) => - get(`/tasks/${id}/asset-refs?asset_id=${assetId}`), + taskAssetRefs: (id: string, assetId: number) => get(`/tasks/${id}/asset-refs?asset_id=${assetId}`), // ---- workspace file manager (workDir) ---- - workspaceList: (path = "") => - get(`/workspace/list?path=${encodeURIComponent(path)}`), - workspaceRead: (path: string) => - get(`/workspace/read?path=${encodeURIComponent(path)}`), + workspaceList: (path = "") => get(`/workspace/list?path=${encodeURIComponent(path)}`), + workspaceRead: (path: string) => get(`/workspace/read?path=${encodeURIComponent(path)}`), workspaceWrite: (path: string, content: string) => post<{ ok: boolean; path: string }>(`/workspace/write`, { path, content }), - workspaceMkdir: (path: string) => - post<{ ok: boolean; path: string }>(`/workspace/mkdir`, { path }), - workspaceDelete: (path: string) => - del<{ ok: boolean }>(`/workspace/delete?path=${encodeURIComponent(path)}`), + workspaceMkdir: (path: string) => post<{ ok: boolean; path: string }>(`/workspace/mkdir`, { path }), + workspaceDelete: (path: string) => del<{ ok: boolean }>(`/workspace/delete?path=${encodeURIComponent(path)}`), workspaceUpload: async (dir: string, files: File[]) => { if (MOCK) return { uploaded: files.length }; const fd = new FormData(); @@ -236,8 +289,9 @@ export const api = { // ---- assets ---- // Server-side paginated: pass limit/offset, get back the page + full match total. assets: (type = "", limit = 50, offset = 0) => - get<{ count: number; total: number; assets: Asset[] }>(`/assets?type=${type}&limit=${limit}&offset=${offset}`) - .then((r) => ({ assets: r?.assets ?? [], total: r?.total ?? r?.count ?? 0 })), + get<{ count: number; total: number; assets: Asset[] }>(`/assets?type=${type}&limit=${limit}&offset=${offset}`).then( + (r) => ({ assets: r?.assets ?? [], total: r?.total ?? r?.count ?? 0 }), + ), searchAssets: (dsl: string, type = "", limit = 50, offset = 0) => get<{ count: number; total: number; assets: Asset[] }>( `/assets?dsl=${encodeURIComponent(dsl)}${type ? `&type=${encodeURIComponent(type)}` : ""}&limit=${limit}&offset=${offset}`, @@ -257,18 +311,38 @@ export const api = { createCompany: (name: string, logo: string, scope: string) => post<{ id: number; created: boolean; scope_added?: number; scope_invalid?: number; scope_errors?: string[] }>( "/companies", - { name, logo, scope: scope.trim() ? scope.split("\n").map((s) => s.trim()).filter(Boolean) : [] }, + { + name, + logo, + scope: scope.trim() + ? scope + .split("\n") + .map((s) => s.trim()) + .filter(Boolean) + : [], + }, ), addCompanyScope: (id: number, scope: string, reason = "") => - post<{ added: number; skipped: number; invalid: number; errors?: string[] }>( - `/companies/${id}/scope`, - { scope: scope.trim() ? scope.split("\n").map((s) => s.trim()).filter(Boolean) : [], reason }, - ), + post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, { + scope: scope.trim() + ? scope + .split("\n") + .map((s) => s.trim()) + .filter(Boolean) + : [], + reason, + }), updateCompanyScope: (id: number, scope: string, reason = "") => - post<{ added: number; skipped: number; invalid: number; errors?: string[] }>( - `/companies/${id}/scope`, - { scope: scope.trim() ? scope.split("\n").map((s) => s.trim()).filter(Boolean) : [], reason, reset: true }, - ), + post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, { + scope: scope.trim() + ? scope + .split("\n") + .map((s) => s.trim()) + .filter(Boolean) + : [], + reason, + reset: true, + }), deleteCompany: (id: number, deleteAssets = false) => http<{ deleted: number; assets_deleted: number }>(`/companies/${id}`, { method: "DELETE", @@ -288,17 +362,52 @@ export const api = { return get(`/exploration/findings?${p.toString()}`); }, findingStats: () => get("/exploration/findings/stats"), + // exportFindings 触发发现页导出并下载文件。scope=selected 时传 ids(finding_id 列表); + // scope=filtered 时传当前筛选(沿用 FindingQuery 的筛选字段);scope=all 忽略筛选。 + exportFindings: async (opts: { + format: "md-single" | "md-zip" | "csv" | "json"; + scope: "filtered" | "all" | "selected"; + filters?: Omit; + ids?: string[]; + }) => { + const p = new URLSearchParams({ format: opts.format, scope: opts.scope }); + if (opts.scope === "selected") { + p.set("ids", (opts.ids ?? []).join(",")); + } else if (opts.scope === "filtered" && opts.filters) { + const f = opts.filters; + if (f.severity && f.severity !== "all") p.set("severity", f.severity); + if (f.status && f.status !== "all") p.set("status", f.status); + if (f.vulnclass && f.vulnclass !== "all") p.set("vulnclass", f.vulnclass); + if (f.task && f.task !== "all") p.set("task_id", f.task); + if (f.sort) p.set("sort", f.sort); + } + const token = getToken(); + const r = await fetch(`/api/exploration/findings/export?${p.toString()}`, { + headers: token ? { Authorization: `Bearer ${token}` } : {}, + }); + if (!r.ok) throw new Error(`export: ${r.status}`); + const blob = await r.blob(); + // 文件名优先取后端 Content-Disposition,取不到则用默认名。 + const disp = r.headers.get("Content-Disposition") ?? ""; + const m = disp.match(/filename="?([^"]+)"?/); + const filename = m?.[1] ?? `findings-export`; + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = filename; + document.body.appendChild(a); + a.click(); + a.remove(); + URL.revokeObjectURL(url); + }, getFinding: (id: string, contextTaskId?: string) => get( `/exploration/findings/${id}${contextTaskId ? `?context_task=${encodeURIComponent(contextTaskId)}` : ""}`, ), // 漏洞链路:该漏洞节点回溯到任务初始节点的子图(节点 + 关系)。 - findingLineage: (id: string) => - get<{ nodes: TaskNode[]; edges: Edge[] }>(`/exploration/findings/${id}/lineage`), - setFindingStatus: (id: string, status: FindingStatus) => - patch(`/exploration/findings/${id}`, { status }), - setFindingSeverity: (id: string, severity: Severity) => - patch(`/exploration/findings/${id}`, { severity }), + findingLineage: (id: string) => get<{ nodes: TaskNode[]; edges: Edge[] }>(`/exploration/findings/${id}/lineage`), + setFindingStatus: (id: string, status: FindingStatus) => patch(`/exploration/findings/${id}`, { status }), + setFindingSeverity: (id: string, severity: Severity) => patch(`/exploration/findings/${id}`, { severity }), // 一次保存漏洞的名称/类别/严重等级(发现列表行内编辑用),只传出现的字段。 updateFinding: ( id: string, @@ -312,10 +421,11 @@ export const api = { workers: arr(r.workers), total: r.total, })), - tokenDaily: (days = 30) => - get(`/tokens/daily?days=${days}`).then(arr), + tokenDaily: (days = 30) => get(`/tokens/daily?days=${days}`).then(arr), conversationTokens: () => - get("/tokens/conversations").then(arr).catch(() => [] as ConvTokenSummary[]), + get("/tokens/conversations") + .then(arr) + .catch(() => [] as ConvTokenSummary[]), explorationGraph: (task?: string) => get<{ nodes: TaskNode[]; edges: Edge[] }>(`/exploration/graph${tq(task)}`), activity: (task?: string, opts?: { intent?: string; since?: number; limit?: number }) => { const q = new URLSearchParams(); @@ -323,7 +433,10 @@ export const api = { if (opts?.intent) q.set("intent", opts.intent); if (opts?.since) q.set("since", String(opts.since)); if (opts?.limit) q.set("limit", String(opts.limit)); - return get<{ items: Activity[]; cursor: number }>(`/exploration/activity?${q.toString()}`).then((r) => ({ items: arr(r.items), cursor: r.cursor ?? 0 })); + return get<{ items: Activity[]; cursor: number }>(`/exploration/activity?${q.toString()}`).then((r) => ({ + items: arr(r.items), + cursor: r.cursor ?? 0, + })); }, activityDetail: (id: number, task?: string) => get<{ detail: string }>(`/exploration/activity/${id}${tq(task)}`), // Reverse-paginated session history. session = "main" | "plan" | "intent:". @@ -366,20 +479,21 @@ export const api = { (method && method !== "all" ? `&method=${encodeURIComponent(method)}` : "") + (q ? `&q=${encodeURIComponent(q)}` : ""), ), - trafficExchange: (id: string) => - get(`/traffic/exchange?id=${encodeURIComponent(id)}`), + trafficExchange: (id: string) => get(`/traffic/exchange?id=${encodeURIComponent(id)}`), trafficHosts: () => get<{ hosts: TrafficHost[] }>(`/traffic/hosts`), - trafficDeleteHost: (host: string) => - del<{ deleted: number }>(`/traffic?host=${encodeURIComponent(host)}`), - trafficDeleteHosts: (hosts: string[]) => - del<{ deleted: number }>(`/traffic/hosts`, { hosts }), + trafficDeleteHost: (host: string) => del<{ deleted: number }>(`/traffic?host=${encodeURIComponent(host)}`), + trafficDeleteHosts: (hosts: string[]) => del<{ deleted: number }>(`/traffic/hosts`, { hosts }), // ---- app settings (runtime toggles) ---- settings: () => get(`/settings`), setSettings: (patch: Partial) => put(`/settings`, patch), // Run a real "test" search with the given (or saved) config to verify it works. - testWebSearch: (patch: { web_search_backend?: string; web_search_proxy?: string; brave_search_api_key?: string; tavily_search_api_key?: string }) => - post<{ ok: boolean; error?: string; count?: number; backend?: string }>(`/settings/web-search/test`, patch), + testWebSearch: (patch: { + web_search_backend?: string; + web_search_proxy?: string; + brave_search_api_key?: string; + tavily_search_api_key?: string; + }) => post<{ ok: boolean; error?: string; count?: number; backend?: string }>(`/settings/web-search/test`, patch), report: async (task?: string) => { if (MOCK) return mockReport(task); const token = getToken(); @@ -393,7 +507,15 @@ export const api = { post<{ reply: string; mode: string }>(`/chat${tq(task)}`, { message, attachments }), // 方式1 文件上传:落到会话/任务工作目录 uploads/，返回可供 agent Read 的相对路径。 chatUpload: async (scope: "task" | "session" | "staging", id: string, files: File[]) => { - if (MOCK) return { attachments: files.map((f) => ({ name: f.name, path: `uploads/${f.name}`, size: f.size, abs: `/mock/${f.name}` })) }; + if (MOCK) + return { + attachments: files.map((f) => ({ + name: f.name, + path: `uploads/${f.name}`, + size: f.size, + abs: `/mock/${f.name}`, + })), + }; const fd = new FormData(); for (const f of files) fd.append("file", f); const token = getToken(); @@ -409,15 +531,37 @@ export const api = { gc: (ttl = 86400) => post<{ removed: number }>(`/gc?ttl=${ttl}`, {}), // ---- LLM ---- - getLLM: () => get<{ configured: boolean; provider: string; model: string; base_url: string; proxy?: string; key_set: boolean; rate_per_second?: number; rate_per_minute?: number; context_window_k?: number; reasoning_effort?: string }>("/llm"), - setLLM: (provider: string, model: string, base_url: string, api_key: string, rate_per_second = 0, rate_per_minute = 0, proxy = "", context_window_k = 0) => - post("/llm", { provider, model, base_url, proxy, api_key, rate_per_second, rate_per_minute, context_window_k }), + getLLM: () => + get<{ + configured: boolean; + provider: string; + model: string; + base_url: string; + proxy?: string; + key_set: boolean; + rate_per_second?: number; + rate_per_minute?: number; + context_window_k?: number; + thinking_type?: string; + reasoning_effort?: string; + }>("/llm"), + setLLM: ( + provider: string, + model: string, + base_url: string, + api_key: string, + rate_per_second = 0, + rate_per_minute = 0, + proxy = "", + context_window_k = 0, + ) => post("/llm", { provider, model, base_url, proxy, api_key, rate_per_second, rate_per_minute, context_window_k }), testLLM: ( provider: string, model: string, base_url: string, api_key: string, proxy = "", + thinking_type = "", reasoning_effort = "", profile_id?: number, ) => @@ -427,6 +571,7 @@ export const api = { base_url, proxy, api_key, + thinking_type, reasoning_effort, profile_id, }), @@ -442,7 +587,8 @@ export const api = { rate_per_second?: number; rate_per_minute?: number; context_window_k?: number; - reasoning_effort?: string; // ""|"off"|"low"|"medium"|"high"|"max" + thinking_type?: string; // ""(不发送)|"disabled"|"enabled" + reasoning_effort?: string; // ""(不发送)|"low"|"medium"|"high"|"xhigh"|"max" priority?: number; // 轮询顺位，越大越先用 pool_exclude?: boolean; // true=不作为故障转移目标 }) => post<{ id: number }>("/llm/profiles", p), @@ -464,8 +610,7 @@ export const api = { // ---- agents ---- agents: () => get<{ agents: Agent[] }>("/agents").then((r) => arr(r.agents)), getAgent: (key: string) => get(`/agents/${key}`), - createAgent: (key: string, name: string, description = "") => - post("/agents", { key, name, description }), + createAgent: (key: string, name: string, description = "") => post("/agents", { key, name, description }), updateAgent: (key: string, name: string, description = "") => patch<{ ok: boolean }>(`/agents/${key}`, { name, description }), deleteAgent: (key: string) => del<{ deleted: string }>(`/agents/${key}`), @@ -497,7 +642,8 @@ export const api = { sendConversationMessage: (id: number, message: string, attachments?: ChatAttachment[]) => post<{ status: string }>(`/conversations/${id}/messages`, { message, attachments }), stopConversation: (id: number) => post<{ status: string }>(`/conversations/${id}/stop`, {}), - saveAgentPrompt: (key: string, template: string, note = "") => put<{ version: number }>(`/agents/${key}/prompt`, { template, note }), + saveAgentPrompt: (key: string, template: string, note = "") => + put<{ version: number }>(`/agents/${key}/prompt`, { template, note }), resetAgentPrompt: (key: string) => post<{ version: number }>(`/agents/${key}/prompt/reset`, {}), // 收尾提示词(超时/步数耗尽的 settlement 提示);prompt 空串=清除覆盖、用内置默认; // max_turns 省略则不动、传 0=用内置默认轮数 @@ -509,7 +655,10 @@ export const api = { saveAgentTaskTimeoutWrapup: (key: string, prompt: string, maxTurns?: number) => put<{ ok: boolean }>(`/agents/${key}/wrapup/task-timeout`, { prompt, max_turns: maxTurns }), resetAgentTaskTimeoutWrapup: (key: string) => - post<{ ok: boolean; task_timeout_wrapup_default: string; task_timeout_wrapup_max_turns_default: number }>(`/agents/${key}/wrapup/task-timeout/reset`, {}), + post<{ ok: boolean; task_timeout_wrapup_default: string; task_timeout_wrapup_max_turns_default: number }>( + `/agents/${key}/wrapup/task-timeout/reset`, + {}, + ), // P3 triggers (仅自定义 agent) agentTriggers: (key: string) => get<{ triggers: AgentTrigger[] }>(`/agents/${key}/triggers`).then((r) => arr(r.triggers)), @@ -531,12 +680,15 @@ export const api = { trigger_max_parallel?: number; }, ) => put<{ ok: boolean }>(`/agents/${key}/config`, patch), - agentPromptVersions: (key: string) => get<{ versions: PromptVersion[] }>(`/agents/${key}/prompts`).then((r) => arr(r.versions)), - agentVariables: (key: string) => get<{ variables: PromptVar[] }>(`/agents/${key}/variables`).then((r) => arr(r.variables)), + agentPromptVersions: (key: string) => + get<{ versions: PromptVersion[] }>(`/agents/${key}/prompts`).then((r) => arr(r.versions)), + agentVariables: (key: string) => + get<{ variables: PromptVar[] }>(`/agents/${key}/variables`).then((r) => arr(r.variables)), previewAgentPrompt: (key: string, template: string, sample?: Record) => post<{ rendered: string; error?: string }>(`/agents/${key}/prompt/preview`, { template, sample }), getAgentVisibility: (key: string) => get<{ mcp: number[]; skill: string[] }>(`/agents/${key}/visibility`), - setAgentVisibility: (key: string, mcp: number[], skill: string[]) => put<{ ok: boolean }>(`/agents/${key}/visibility`, { mcp, skill }), + setAgentVisibility: (key: string, mcp: number[], skill: string[]) => + put<{ ok: boolean }>(`/agents/${key}/visibility`, { mcp, skill }), // ---- tools (内置工具目录) ---- tools: () => get<{ tools: Tool[] }>("/tools").then((r) => arr(r.tools)), @@ -544,10 +696,13 @@ export const api = { put<{ ok: boolean }>(`/tools/${key}`, patch), resetTool: (key: string) => post<{ ok: boolean }>(`/tools/${key}/reset`, {}), // custom tools (自定义工具) - createCustomTool: (t: Pick) => - post<{ key: string }>("/tools/custom", t), - updateCustomTool: (key: string, t: Pick) => - put<{ ok: boolean }>(`/tools/custom/${key}`, t), + createCustomTool: ( + t: Pick, + ) => post<{ key: string }>("/tools/custom", t), + updateCustomTool: ( + key: string, + t: Pick, + ) => put<{ ok: boolean }>(`/tools/custom/${key}`, t), deleteCustomTool: (key: string) => del<{ deleted: string }>(`/tools/custom/${key}`), testCustomTool: (body: { kind: string; exec: Record; params: Record }) => post<{ output: string; is_error: boolean }>("/tools/custom/test", body), @@ -591,8 +746,14 @@ export const api = { // ---- skills (文件系统) ---- skills: () => get<{ skills: SkillItem[] }>("/skills").then((r) => arr(r.skills)), - createSkill: (s: { name: string; description: string; license?: string; compatibility?: string; mcps?: string[]; instructions?: string }) => - post<{ name: string }>("/skills", s), + createSkill: (s: { + name: string; + description: string; + license?: string; + compatibility?: string; + mcps?: string[]; + instructions?: string; + }) => post<{ name: string }>("/skills", s), // uploadSkill installs a skill from a .zip (multipart). Surfaces the backend // error text (e.g. 已存在 / 缺少 SKILL.md) so the UI can show a precise message. uploadSkill: async (file: File, overwrite = false): Promise<{ name: string; files: number }> => { @@ -610,20 +771,25 @@ export const api = { return body; }, deleteSkill: (name: string) => del<{ deleted: string }>(`/skills/${name}`), - updateSkillMeta: (name: string, meta: { mcps?: string[]; description?: string; license?: string; compatibility?: string }) => - put<{ ok: boolean }>(`/skills/${name}/meta`, meta), - createSkillDir: (skill: string, path: string) => - post<{ dir: string }>(`/skills/${skill}/dirs`, { path }), + updateSkillMeta: ( + name: string, + meta: { mcps?: string[]; description?: string; license?: string; compatibility?: string }, + ) => put<{ ok: boolean }>(`/skills/${name}/meta`, meta), + createSkillDir: (skill: string, path: string) => post<{ dir: string }>(`/skills/${skill}/dirs`, { path }), skillFiles: (name: string) => get<{ files: string[] }>(`/skills/${name}/files`).then((r) => r.files), readSkillFile: (name: string, file: string) => get<{ content: string; file: string }>(`/skills/${name}/files/${file}`).then((r) => r.content), writeSkillFile: (name: string, file: string, content: string) => put<{ ok: boolean }>(`/skills/${name}/files/${file}`, { content }), - deleteSkillPath: (skill: string, path: string) => - del<{ deleted: string }>(`/skills/${skill}/files/${path}`), + deleteSkillPath: (skill: string, path: string) => del<{ deleted: string }>(`/skills/${skill}/files/${path}`), + skillUsage: (name: string, limit = 50) => + get<{ calls: SkillCall[] }>(`/skills/${name}/usage?limit=${limit}`).then((r) => arr(r.calls)), + missingSkills: (limit = 20) => + get<{ missing: MissingSkill[] }>(`/skills/missing?limit=${limit}`).then((r) => arr(r.missing)), // ---- visibility (MCP resource side) ---- (agent ids are strings per spec) - resourceVisibility: (kind: string, id: number) => get<{ agents: string[] }>(`/visibility/${kind}/${id}`).then((r) => arr(r.agents)), + resourceVisibility: (kind: string, id: number) => + get<{ agents: string[] }>(`/visibility/${kind}/${id}`).then((r) => arr(r.agents)), toggleVisibility: (agentId: string, kind: string, resourceId: number, visible: boolean) => post<{ ok: boolean }>("/visibility/toggle", { agent_id: agentId, kind, resource_id: resourceId, visible }), @@ -648,7 +814,8 @@ export const api = { interceptDecide: (id: number, decision: "allowed" | "denied") => post<{ ok: boolean }>(`/intercept/pending/${id}/decide`, { decision }), interceptHistory: () => get<{ items: InterceptApprovalRow[] }>("/intercept/history").then((r) => arr(r.items)), - interceptTask: (taskId: string) => get<{ items: InterceptApprovalRow[] }>(`/intercept/task/${taskId}`).then((r) => arr(r.items)), + interceptTask: (taskId: string) => + get<{ items: InterceptApprovalRow[] }>(`/intercept/task/${taskId}`).then((r) => arr(r.items)), // ---- intercept tool-config (全局工具拦截范围) ---- interceptGetToolConfig: async (): Promise<{ enabled_tools: string[] }> => { @@ -696,6 +863,9 @@ export const api = { }, llmRecordDetail: (id: number) => get(`/llm/records/${id}`), llmTasks: () => get<{ tasks: LLMTask[] }>(`/llm/records/tasks`), - llmRecordsDeleteTask: (task: string) => - del<{ deleted: number }>(`/llm/records?task=${encodeURIComponent(task)}`), + llmRecordsDeleteTask: (task: string) => del<{ deleted: number }>(`/llm/records?task=${encodeURIComponent(task)}`), + // 按模型聚合本任务的 token 用量（来自常开的 llm_usage 计量账本，逐次调用精确， + // per-agent 绑定 / 轮询 / 中断消耗都覆盖）。 + tokensByModel: (task: string) => + get<{ models: ModelTokenStat[] }>(`/llm/records/by-model?task=${encodeURIComponent(task)}`), }; diff --git a/web/src/lib/mock/data.ts b/web/src/lib/mock/data.ts index 044457f..ae772e7 100644 --- a/web/src/lib/mock/data.ts +++ b/web/src/lib/mock/data.ts @@ -7,7 +7,7 @@ import type { Activity, Agent, AgentDetail, Asset, Audit, Company, Conversation, ConvTokenSummary, DailyTokenBucket, Edge, Finding, InterceptApprovalRow, InterceptPending, InterceptRule, LLMPoolStatus, LLMProfile, MCPServer, MCPTool, PromptVar, - PromptVersion, Settings, SkillItem, Stats, TaskNode, Task, TokenTotal, + PromptVersion, Settings, SkillCall, SkillItem, MissingSkill, Stats, TaskNode, Task, TokenTotal, TokenUsage, Tool, TrafficResp, TrafficDetail, TrafficHost, LLMTask, } from "@/lib/types"; @@ -617,14 +617,14 @@ export const audit: Audit = { // ── LLM profiles ───────────────────────────────────────────────────────────── export const llmProfiles: LLMProfile[] = [ - { id: "1", name: "Claude Opus 4.8", format: "anthropic", model: "claude-opus-4-8", api_key_hint: "…a3f2", rate_per_second: 0, rate_per_minute: 0, context_window_k: 1000, reasoning_effort: "high", is_default: true, priority: 0, pool_exclude: false }, - { id: "2", name: "DeepSeek V4", format: "openai", base_url: "https://api.deepseek.com", model: "deepseek-v4-flash", api_key_hint: "…9c11", rate_per_second: 0, rate_per_minute: 60, context_window_k: 128, reasoning_effort: "", is_default: false, priority: 10, pool_exclude: false }, + { id: "1", name: "Claude Opus 4.8", format: "anthropic", model: "claude-opus-4-8", api_key_hint: "…a3f2", rate_per_second: 0, rate_per_minute: 0, context_window_k: 1000, thinking_type: "enabled", reasoning_effort: "high", is_default: true, priority: 0, pool_exclude: false }, + { id: "2", name: "DeepSeek V4", format: "openai", base_url: "https://api.deepseek.com", model: "deepseek-v4-flash", api_key_hint: "…9c11", rate_per_second: 0, rate_per_minute: 60, context_window_k: 128, thinking_type: "", reasoning_effort: "", is_default: false, priority: 10, pool_exclude: false }, ]; export const llmConfig = { configured: true, provider: "anthropic", model: "claude-opus-4-8", base_url: "", proxy: "", key_set: true, - rate_per_second: 0, rate_per_minute: 0, context_window_k: 1000, reasoning_effort: "high", + rate_per_second: 0, rate_per_minute: 0, context_window_k: 1000, thinking_type: "enabled", reasoning_effort: "high", }; // ── Agents ─────────────────────────────────────────────────────────────────── @@ -686,9 +686,19 @@ export const mcpToolsById: Record = { // ── Skills ─────────────────────────────────────────────────────────────────── export const skills: SkillItem[] = [ - { name: "api-recon", description: "对 REST/GraphQL API 做侦察与越权面枚举；发现新 API 端点时使用。", license: "MIT", mcps: [], files: ["SKILL.md", "scripts/enum.py"] }, - { name: "playwright-cli", description: "用 Playwright 驱动浏览器做动态爬取与截图；需要渲染 JS 站点时使用。", mcps: ["playwright"], files: ["SKILL.md"] }, - { name: "scopesentry", description: "从 ScopeSentry 拉取资产并归并到公司范围；批量导入资产时使用。", files: ["SKILL.md", "assets/mapping.md"] }, + { name: "api-recon", description: "对 REST/GraphQL API 做侦察与越权面枚举；发现新 API 端点时使用。", license: "MIT", mcps: [], files: ["SKILL.md", "scripts/enum.py"], calls: 24, tasks: 6, usage_agents: ["worker", "planner"], last_used: "2026-08-20T09:12:00Z" }, + { name: "playwright-cli", description: "用 Playwright 驱动浏览器做动态爬取与截图；需要渲染 JS 站点时使用。", mcps: ["playwright"], files: ["SKILL.md"], calls: 7, tasks: 3, usage_agents: ["worker"], last_used: "2026-08-19T21:40:00Z" }, + { name: "scopesentry", description: "从 ScopeSentry 拉取资产并归并到公司范围；批量导入资产时使用。", files: ["SKILL.md", "assets/mapping.md"], calls: 0, tasks: 0, usage_agents: [] }, +]; + +export const skillCalls: SkillCall[] = [ + { ts: "2026-08-20T09:12:00Z", agent_key: "worker", task_id: 42, session_id: "", args_len: 128 }, + { ts: "2026-08-20T08:03:00Z", agent_key: "planner", task_id: 42, session_id: "", args_len: 0 }, + { ts: "2026-08-19T17:55:00Z", agent_key: "worker", task_id: 37, session_id: "", args_len: 64 }, +]; + +export const missingSkills: MissingSkill[] = [ + { skill: "jwt-forge", calls: 3, agents: ["worker"], last_used: "2026-08-20T07:20:00Z" }, ]; // ── Tools ──────────────────────────────────────────────────────────────────── diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 1b57783..d534ed7 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -149,6 +149,74 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (seg[0] === "tasks" && seg[2] === "coverage-graph") return D.coverageGraph; if (seg[0] === "tasks" && seg[2] === "asset-refs") return D.assetRefsFor(Number(q.get("asset_id") ?? 0)); + // ── 任务测试范围（增删查）── + if (seg[0] === "tasks" && seg[2] === "scope" && seg.length === 3 && m === "GET") return { scope: [] }; + if (seg[0] === "tasks" && seg[2] === "scope" && seg.length === 3 && m === "POST") + return { id: Date.now(), task_id: Number(seg[1]), kind: b.kind, domain: b.value, source: "manual" }; + if (seg[0] === "tasks" && seg[2] === "scope" && seg.length === 4 && m === "DELETE") return { ok: true }; + + // ── 全局 llm_usage 聚合（仪表盘新版视图，demo）── + if (path === "/tokens/usage") + return { + by_profile: [ + { + profile_name: "default", + calls: 60, + tasks: 4, + input_tokens: 1570000, + output_tokens: 110000, + cache_read_tokens: 1120000, + cache_write_tokens: 140000, + }, + ], + daily: [ + { + profile_name: "default", + date: "2026-08-18", + input_tokens: 520000, + output_tokens: 38000, + cache_read_tokens: 370000, + }, + { + profile_name: "default", + date: "2026-08-19", + input_tokens: 640000, + output_tokens: 45000, + cache_read_tokens: 460000, + }, + { + profile_name: "default", + date: "2026-08-20", + input_tokens: 410000, + output_tokens: 27000, + cache_read_tokens: 290000, + }, + ], + }; + + // ── 按模型 token 用量（demo：一条示例）── + if (path === "/llm/records/by-model") + return { + models: [ + { + model: "claude-opus-4-6", + calls: 42, + input_tokens: 1250000, + output_tokens: 86000, + cache_read_tokens: 940000, + cache_write_tokens: 120000, + }, + { + model: "claude-haiku-4-5", + calls: 18, + input_tokens: 320000, + output_tokens: 24000, + cache_read_tokens: 180000, + cache_write_tokens: 20000, + }, + ], + }; + // ── 工作空间文件管理器（demo：静态示例树；写/建/删走下方写兜底 {ok:true}）── if (path === "/workspace/list") return D.workspaceList(q.get("path") ?? ""); if (path === "/workspace/read") return D.workspaceRead(q.get("path") ?? ""); @@ -200,18 +268,12 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re }; } // 单条 finding:GET 详情 / PATCH 改状态/严重度/名称/类别(demo 直接改内存对象)。 - if ( - seg[0] === "exploration" && - seg[1] === "findings" && - seg.length === 3 && - seg[2] !== "stats" - ) { + if (seg[0] === "exploration" && seg[1] === "findings" && seg.length === 3 && seg[2] !== "stats") { const f = mockFindings.find((x) => x.id === seg[2]); if (!f) return {}; if (m === "PATCH") { if (typeof b.status === "string") f.status = b.status as typeof f.status; - if (typeof b.severity === "string") - f.severity = b.severity as typeof f.severity; + if (typeof b.severity === "string") f.severity = b.severity as typeof f.severity; if (typeof b.name === "string") f.name = b.name; if (typeof b.vulnclass === "string") f.vulnclass = b.vulnclass; } @@ -288,7 +350,9 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/exploration/activity" && seg.length === 2) { const since = Number(q.get("since") ?? 0); const limit = Math.max(1, Number(q.get("limit") ?? 300)); - const items = D.activityForTask().filter((item) => item.seq > since).slice(0, limit); + const items = D.activityForTask() + .filter((item) => item.seq > since) + .slice(0, limit); return { items, cursor: items.length ? items[items.length - 1].seq : since }; } if (seg[0] === "exploration" && seg[1] === "activity" && seg.length === 3) { @@ -309,7 +373,8 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/settings" && m === "PUT") return { ...D.settings, ...b }; if (path === "/settings/web-search/test") return { ok: true, count: 5, backend: D.settings.web_search_backend }; if (path === "/settings/python/detect") return { python_interpreter: "/usr/bin/python3" }; - if (path === "/chat") return { reply: "（demo）我已把该建议注入为一条高优意图，work agent 会尽快执行。", mode: "hint" }; + if (path === "/chat") + return { reply: "（demo）我已把该建议注入为一条高优意图，work agent 会尽快执行。", mode: "hint" }; if (path === "/gc") return { removed: 0 }; // ── 工具执行历史 ── @@ -337,11 +402,20 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/llm/pool" && m === "GET") return D.llmPool; if (path === "/llm/pool/reset") return { ...D.llmPool, chain: D.llmPool.chain.map((c) => ({ ...c, state: "ok", fails: 0, cooldown_secs: 0 })) }; - if (seg[0] === "llm" && seg[1] === "profiles" && seg.length === 3 && m === "DELETE") return { deleted: Number(seg[2]) }; + if (seg[0] === "llm" && seg[1] === "profiles" && seg.length === 3 && m === "DELETE") + return { deleted: Number(seg[2]) }; // ── agents ── if (path === "/agents" && m === "GET") return { agents: D.agents }; - if (path === "/agents" && m === "POST") return { id: "9", key: String(b.key ?? "custom"), name: String(b.name ?? ""), role: "custom", builtin: false, enabled: true }; + if (path === "/agents" && m === "POST") + return { + id: "9", + key: String(b.key ?? "custom"), + name: String(b.name ?? ""), + role: "custom", + builtin: false, + enabled: true, + }; if (seg[0] === "agents" && seg.length === 2 && m === "GET") return D.agentDetail(seg[1]); if (seg[0] === "agents" && seg[2] === "triggers" && m === "GET") return { triggers: [] }; if (seg[0] === "agents" && seg[2] === "prompts") return { versions: D.agentDetail(seg[1]).versions }; @@ -352,7 +426,14 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re // ── conversations ── if (path === "/conversations" && m === "GET") return { conversations: D.conversations }; - if (path === "/conversations" && m === "POST") return { id: 3, agent_key: String(b.agent_key ?? "mainagent"), title: String(b.title ?? "新会话"), created_at: "2026-07-26T04:00:00Z", updated_at: "2026-07-26T04:00:00Z" }; + if (path === "/conversations" && m === "POST") + return { + id: 3, + agent_key: String(b.agent_key ?? "mainagent"), + title: String(b.title ?? "新会话"), + created_at: "2026-07-26T04:00:00Z", + updated_at: "2026-07-26T04:00:00Z", + }; if (seg[0] === "conversations" && seg[2] === "messages" && seg.length === 3 && m === "GET") { const items = D.conversationMessages[Number(seg[1])] ?? []; return { items, cursor: items.length ? items[items.length - 1].seq : 0, running: false }; @@ -378,13 +459,16 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (seg[0] === "mcp" && seg.length === 2 && m === "DELETE") return { deleted: Number(seg[1]) }; // ── scopesentry（demo：未配置）── - if (path === "/sync/scopesentry/status") return { exists: false, configured: false, enabled: false, reachable: false, tools: [] }; + if (path === "/sync/scopesentry/status") + return { exists: false, configured: false, enabled: false, reachable: false, tools: [] }; if (path === "/sync/scopesentry/projects") return { projects: [], tag: {} }; if (path === "/sync/scopesentry/tasks") return { tasks: [] }; if (path === "/sync/scopesentry/sync") return { synced: {}, companies: null, warnings: null, errors: null }; // ── skills ── if (path === "/skills" && m === "GET") return { skills: D.skills }; + if (path === "/skills/missing") return { missing: D.missingSkills }; + if (seg[0] === "skills" && seg[2] === "usage") return { calls: D.skillCalls }; if (path === "/skills" && m === "POST") return { name: String(b.name ?? "new-skill") }; if (seg[0] === "skills" && seg[2] === "files" && seg.length === 3) return { files: ["SKILL.md"] }; if (seg[0] === "skills" && seg[2] === "files" && seg.length >= 4) @@ -395,7 +479,8 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re // ── intercept ── if (path === "/intercept/rules" && m === "GET") return { rules: D.interceptRules }; - if (seg[0] === "intercept" && seg[1] === "rules" && seg[3] === "toggle") return { ok: true, enabled: b.enabled ?? true }; + if (seg[0] === "intercept" && seg[1] === "rules" && seg[3] === "toggle") + return { ok: true, enabled: b.enabled ?? true }; if (path === "/intercept/pending" && m === "GET") return { pending: D.interceptPending }; if (seg[0] === "intercept" && seg[1] === "pending" && seg[3] === "decide") return { ok: true }; if (seg[0] === "intercept" && seg[1] === "pending" && seg.length === 3 && m === "GET") @@ -409,5 +494,9 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (["POST", "PUT", "PATCH", "DELETE"].includes(m)) return { ok: true }; // ── 读兜底：集合类给 []，其余 {} ── - return /(\/(tasks|profiles|conversations|rules|history|projects|tokens|agents|servers|skills|tools|findings|intents)s?$)|s$/.test(path) ? [] : {}; + return /(\/(tasks|profiles|conversations|rules|history|projects|tokens|agents|servers|skills|tools|findings|intents)s?$)|s$/.test( + path, + ) + ? [] + : {}; } diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index ca761f7..141aeee 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -225,14 +225,26 @@ export interface WorkspaceFile { content?: string; } +// 任务测试范围的一条（覆盖度分母 + 授权边界）。 +export interface TaskScopeRow { + id: number; + task_id: number; + kind: "company" | "root_domain" | "subdomain" | "ip" | "cidr"; + company_id?: number; + domain?: string; + net?: string; + source: "auto" | "agent" | "manual"; + reason?: string; +} + // 公司资产范围规则的一条（归属唯一真值来源）。 export interface ScopeRow { id: number; company_id: number; kind: "domain" | "ip" | "cidr"; - domain?: string; // kind=domain 时有值 - net?: string; // kind=ip|cidr 时有值 - raw: string; // 原始用户输入，用于显示和回填 + domain?: string; // kind=domain 时有值 + net?: string; // kind=ip|cidr 时有值 + raw: string; // 原始用户输入，用于显示和回填 reason?: string; } @@ -353,11 +365,11 @@ export type ActivityKind = | "thinking" | "result" | "user" - | "intent" // LLM-generated exploration objective leading a worker session (UI-synthesized) - | "round" // planner round boundary marker (engine-emitted) - | "usage" // live cumulative token usage (per model turn); not rendered - | "llm_switch" // automatic/manual task-level LLM switch - | "llm_failover" // task-level provider switch / chain exhaustion audit event + | "intent" // LLM-generated exploration objective leading a worker session (UI-synthesized) + | "round" // planner round boundary marker (engine-emitted) + | "usage" // live cumulative token usage (per model turn); not rendered + | "llm_switch" // automatic/manual task-level LLM switch + | "llm_failover" // task-level provider switch / chain exhaustion audit event | "intercept_request"; // user-approval request from the intercept layer // ChatAttachment 是一次上传的文件:path 相对该会话/任务工作目录(即 agent 的 CWD)。 @@ -475,6 +487,43 @@ export interface TokenTotal { cache_write_tokens: number; } +// Global per-profile token spend from the llm_usage ledger (GET /api/tokens/usage). +export interface ProfileUsage { + profile_name: string; + calls: number; + tasks: number; + input_tokens: number; + output_tokens: number; + cache_read_tokens: number; + cache_write_tokens: number; +} + +// One (profile, UTC day) token bucket for the dashboard's daily chart (new source). +export interface ProfileDayUsage { + profile_name: string; + date: string; // YYYY-MM-DD + input_tokens: number; + output_tokens: number; + cache_read_tokens: number; +} + +// Response of GET /api/tokens/usage — the dashboard's "new" (llm_usage) token view. +export interface UsageStats { + by_profile: ProfileUsage[]; + daily: ProfileDayUsage[]; +} + +// Per-model token usage for one task (GET /api/llm/records/by-model), from the +// always-on llm_usage metering ledger. calls = number of LLM calls on this model. +export interface ModelTokenStat { + model: string; + calls: number; + input_tokens: number; + output_tokens: number; + cache_read_tokens: number; + cache_write_tokens: number; +} + export interface Session { id: string; role: SessionRole; @@ -574,7 +623,9 @@ export interface LLMProfile { rate_per_second: number; rate_per_minute: number; context_window_k?: number; - // 思考模式: ""=默认(不发送) | "off"=关闭 | "low"/"medium"/"high"/"max"=开启并设强度 + // 思考开关(thinking.type): ""=不发送(默认) | "disabled"=关闭 | "enabled"=开启 + thinking_type?: string; + // 思考强度: ""=不发送(默认) | "low"/"medium"/"high"/"xhigh"/"max" reasoning_effort?: string; is_default: boolean; // 轮询顺位：越大越先被选中。激活配置恒为链首，与本值无关。 @@ -690,12 +741,34 @@ export interface MCPTool { // Fields align with the agentskills.io open specification. // description covers both "what the skill does" and "when to use it". export interface SkillItem { - name: string; // unique key = directory name - description?: string; // required per spec; covers what + when to use - license?: string; // optional: SPDX identifier or free text + name: string; // unique key = directory name + description?: string; // required per spec; covers what + when to use + license?: string; // optional: SPDX identifier or free text compatibility?: string; // optional: environment requirements - mcps?: string[]; // MCP server names this skill unlocks on load - files: string[]; // files in the skill directory + mcps?: string[]; // MCP server names this skill unlocks on load + files: string[]; // files in the skill directory + // 调用统计（skill_usage 账本）。从未被调用过的 skill：calls=0、last_used 缺省。 + calls: number; + tasks: number; // 加载过它的任务数（chat 会话不计入） + usage_agents: string[]; // 加载过它的 agent key + last_used?: string; +} + +// SkillCall 是一次 Skill() 调用（单个 skill 的最近调用列表）。 +export interface SkillCall { + ts: string; + agent_key: string; + task_id: number; // 0 = 非任务场景（对话会话） + session_id: string; + args_len: number; +} + +// MissingSkill 是被点名但不存在的 skill —— "想用但没有"的缺口。 +export interface MissingSkill { + skill: string; + calls: number; + agents: string[]; + last_used?: string; } // ---- Tools (内置工具目录) ---- @@ -708,12 +781,12 @@ export interface Tool { description: string; // eslint-disable-next-line @typescript-eslint/no-explicit-any schema: Record; // full JSON-Schema (object with properties) - agents: string[]; // bound agent keys + agents: string[]; // bound agent keys enabled: boolean; kind?: "builtin" | "shell" | "command" | "script" | "http"; // 自定义工具类型 // eslint-disable-next-line @typescript-eslint/no-explicit-any - exec?: Record; // 自定义工具执行规格(kind!=builtin) - deferred?: boolean; // schema 延迟(SearchExtraTools/ExecuteExtraTool) + exec?: Record; // 自定义工具执行规格(kind!=builtin) + deferred?: boolean; // schema 延迟(SearchExtraTools/ExecuteExtraTool) } // ---- Stats ---- @@ -764,14 +837,14 @@ export interface InterceptPending { // InterceptApprovalRow enriches InterceptPending with conversation/task and rule context. export interface InterceptApprovalRow extends InterceptPending { - conv_title: string; // "" if no linked conversation - conv_agent_key: string; // "" if no linked conversation - rule_name: string; // "" if rule was deleted + conv_title: string; // "" if no linked conversation + conv_agent_key: string; // "" if no linked conversation + rule_name: string; // "" if rule was deleted } // ── 资产同步 (ScopeSentry 数据源) ────────────────────────────────────────────── export interface SSProject { - id: string; // MongoDB ObjectID — used as filter.project + id: string; // MongoDB ObjectID — used as filter.project name: string; logo?: string; AssetCount?: number; @@ -780,7 +853,7 @@ export interface SSProject { export interface SSTask { id: string; - name: string; // used as filter.task + name: string; // used as filter.task status?: number; progress?: number; creatTime?: string; diff --git a/web/src/lib/utils.ts b/web/src/lib/utils.ts index eb6d2c9..e5cbf1d 100644 --- a/web/src/lib/utils.ts +++ b/web/src/lib/utils.ts @@ -45,6 +45,35 @@ export function radixOverlayWasOpenAtPointerDown(): boolean { return overlayOpenAtLastPointerDown; } +// copyText 把文本写入剪贴板,返回是否成功。 +// 背景:navigator.clipboard 仅在安全上下文(HTTPS / localhost)可用;通过 IP + HTTP +// 访问时它为 undefined,此时降级到 execCommand("copy")。 +export async function copyText(text: string): Promise { + if (navigator.clipboard && window.isSecureContext) { + try { + await navigator.clipboard.writeText(text); + return true; + } catch { + // 继续走降级方案 + } + } + try { + const textarea = document.createElement("textarea"); + textarea.value = text; + textarea.style.position = "fixed"; + textarea.style.left = "-9999px"; + textarea.style.top = "0"; + document.body.appendChild(textarea); + textarea.focus(); + textarea.select(); + const ok = document.execCommand("copy"); + document.body.removeChild(textarea); + return ok; + } catch { + return false; + } +} + export const getInitials = (str: string): string => { if (typeof str !== "string" || !str.trim()) return "?"; From d3183dbbd280cb6de4c48aaa8d38a37d015c238b Mon Sep 17 00:00:00 2001 From: pphua Date: Fri, 21 Aug 2026 21:01:39 +0800 Subject: [PATCH 2/8] 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 && ( <> )}