Merge PR #42: 0.3.3 Worker/LLM/Agent/UI 增强

来自 @neouks 的 PR #42,外加本地的三个后续修改:

- fix(llm): Agent 绑定优先于任务 LLM 配置链
- feat(ui): 会话徽标改显示 LLM 配置名
- fix(scope): 含点号的文本不再被误判为 ICP 备案号

已知遗留(code review 确认,评估影响可控,后续单独处理):
- db/exploration.go CancelIntent 先数后删的并发竞态,会回滚并向用户报
  "intent cleanup changed concurrently",不损坏数据
- db/assets.go company_source 优先级 CASE 在 8 处手写重复,改动易漏
- server/goals.go reconcileConcurrency 第二个 if 是死代码

Claude-Session: https://claude.ai/code/session_01Q8nhYZfS4H8fsSCBJXxbSy
This commit is contained in:
Autumn-27 committed 2026-08-22 06:26:43 -04:00
commit f0c2b43802
78 files changed
+12765 -2651

No files matched your search

+23 -20
View File
@@ -48,35 +48,33 @@ jobs:
with:
name: web-dist
path: server/webui/dist
- name: Build
- name: Build and package with build.sh
env:
GOOS: ${{ matrix.goos }}
GOARCH: ${{ matrix.goarch }}
CGO_ENABLED: "0"
ARTEX_TARGET_OS: ${{ matrix.goos }}
ARTEX_TARGET_ARCH: ${{ matrix.goarch }}
ARTEX_BUILD_VERSION: ${{ github.ref_name }}
ARTEX_SKIP_FRONTEND: "1"
ARTEX_SKIP_NPM_CI: "1"
# UPX 自解压 ELF 在部分 Linux 内核、虚拟化和安全策略下会崩溃。
# Release 使用 Go linker 裁剪和 zip 压缩,优先保证可移植性。
ARTEX_COMPRESS: "0"
ARTEX_PACKAGE: "1"
ARTEX_PACKAGE_DIR: dist
run: |
BIN=artex
[ "$GOOS" = "windows" ] && BIN=artex.exe
go build -tags embedui -trimpath \
-ldflags "-s -w -X main.version=${GITHUB_REF_NAME}" \
-o "$BIN" ./cmd/artex
- name: Package
run: |
NAME="artex-${GITHUB_REF_NAME}-${{ matrix.goos }}-${{ matrix.goarch }}"
mkdir -p "$NAME"
cp -r skills "$NAME"/
cp config.example.json "$NAME"/
if [ "${{ matrix.goos }}" = "windows" ]; then cp artex.exe "$NAME"/; else cp artex "$NAME"/; fi
zip -r "$NAME.zip" "$NAME"
./build.sh --target "${ARTEX_TARGET_OS}/${ARTEX_TARGET_ARCH}"
- name: Smoke test Linux amd64 binary
if: matrix.goos == 'linux' && matrix.goarch == 'amd64'
run: dist/artex-linux-amd64/artex -h
- uses: actions/upload-artifact@v4
with:
name: pkg-${{ matrix.goos }}-${{ matrix.goarch }}
path: "*.zip"
path: "dist/*.zip"
# Linux 生二进制单独上传,供 docker job 复用(免在镜像里重编前端+Go)
- if: matrix.goos == 'linux'
uses: actions/upload-artifact@v4
with:
name: bin-linux-${{ matrix.goarch }}
path: artex
path: dist/artex-linux-${{ matrix.goarch }}/artex
# ③ 汇总所有 zip → 创建 GitHub Release
release:
@@ -88,9 +86,14 @@ jobs:
pattern: pkg-*
merge-multiple: true
path: dist
- name: Generate checksums
working-directory: dist
run: sha256sum *.zip > SHA256SUMS
- uses: softprops/action-gh-release@v2
with:
files: dist/*.zip
files: |
dist/*.zip
dist/SHA256SUMS
generate_release_notes: true
# ④ 多架构镜像 → 推 Docker Hub(复用 binaries 编好的 Linux 二进制,
+96 -1
View File
@@ -2,6 +2,100 @@
本项目的重要变更记录在此文件中,格式参考 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/)。
## [Unreleased]
### Build
- `build.sh --release` 支持一次构建 Linux amd64/arm64、macOS amd64/arm64 和 Windows amd64,并生成带 `skills/`、配置示例和 README 的 zip 发布包。
- 使用 Go linker 去除调试信息并生成 zip 发布包;UPX 改为显式可选,避免其自解压 ELF 在部分 Linux 环境启动时发生段错误。
- Release Workflow 对 Linux amd64 二进制执行启动冒烟测试,并随发布包提供 `SHA256SUMS`。
## [0.3.3] - 2026-08-21
### Worker
#### 新增的功能
- 运行中的 Worker 支持单项暂停、恢复和取消;暂停保留意图、事实与漏洞,取消会在执行退出后事务性清理当前意图及其直接产物。
- 增加 `paused` 意图状态、执行栅栏和具名终止原因,防止暂停、取消或任务删除后出现迟到黑板写入。
- 增加可选任务并发上限;新建、恢复、漏洞深入和队列补位统一通过持久化 FIFO 准入路径。
#### 修改的功能
- Worker 单次运行墙钟默认时长由 600 秒调整为 1200 秒;超时任务复活后在下一次真实执行时重新计时。
- Worker、Planner 和主 Agent 的真实执行状态统一驱动任务状态徽标,修复 Worker 运行时仍显示“空闲”的问题。
- Worker 会话保留单项控制和当前会话 Token 汇总,模型名称移动到当前会话标题右侧展示。
#### 删除的功能
- 删除 Worker 多选、全选及批量暂停/继续界面,同时移除 Worker 批量控制 API 和 Mock 契约。
- 删除 Worker 列表行中的 Token 徽标与 Tooltip,底层完整 Token 账本和任务聚合接口继续保留。
### LLM
#### 新增的功能
- 任务支持有序 LLM 配置链;明确识别供应商额度不足时自动切换下一配置,并持久化当前配置、耗尽状态和错误摘要。
- 运行中和暂停任务支持编辑完整配置链、调整顺序和手动切换当前配置,从下一次 LLM 调用生效。
- 自动切换、手动切换和全链耗尽写入结构化系统活动,并通过任务活动流弹出可去重提醒。
- 增加任务角色模型解析接口,统一解析 Main Agent、Planner 和 Worker 下一次调用使用的配置与模型。
#### 修改的功能
- 目标拆解、Planner、Worker 和主 Agent 共享任务级 LLM runtime;任务显式配置链优先于角色绑定和全局配置。
- 故障转移只响应明确的额度、余额或账单错误;普通限流、鉴权、网络、服务端和上下文错误不会错误切换。
- LLM 配置页改为配置卡片列表和右侧抽屉编辑,保留任务级配置链,不引入 Pool 轮询、优先级和熔断界面。
- 配置链耗尽信息支持窄屏换行;当前模型由会话列表图标改为当前会话标题旁的文字标签和完整 Tooltip。
#### 删除的功能
- 删除会话列表中的 LLM 图标式标识,避免与 Worker 状态图标混淆。
### UI
#### 新增的功能
- 增加全局任务模板 CRUD 和右侧管理抽屉;新建任务可加载预设描述与目标,也可将当前内容另存为模板。
- 新建任务支持关联多个来源任务和多个企业资产范围;企业范围通过单个多行文本框自动识别域名、URL、IP、CIDR、ICP 和企业关键词。
- 发现页按任务分组并提供组内独立分页,支持为漏洞填写描述并创建高优先级深入意图。
- 对话支持重命名、置顶、取消置顶和删除;任务列表支持当前页批量暂停与继续。
- 流量详情和工具执行详情改为右侧抽屉;HTTP 报文补齐 `Host` 并高亮请求行、状态码、Header、JSON 和标记语言正文。
#### 修改的功能
- 任务标题改为可聚焦详情链接;任务资产、企业资产、发现任务组和组内漏洞使用服务端真实分页、稳定排序与精确总数。
- 新增企业改用右侧抽屉,范围录入统一为多行文本并提供实时识别、校验和类型预览。
- 主 Agent 输入框支持自动增长、`Enter` 发送、`Shift+Enter` 换行,并避免中文输入法组合态误发送。
- Agent 预览、任务报告和相关详情统一使用共享 Markdown 渲染;删除确认、长错误和移动端抽屉宽度统一修复。
- 任务模板选择框以模板名称显示和搜索、以模板 ID 提交;恢复系统原始字体尺寸,应用版本统一显示为 `0.3.3`。
- 任务总计和当前会话 Token 只高亮输入、缓存读取和输出数值;发送按钮使用更简洁的上箭头图标。
#### 删除的功能
- 删除任务卡片独立“进入”按钮,统一点击任务标题进入详情。
- 删除仪表盘“新建任务”按钮、新增企业的 Logo URL 输入项,以及发现页顶部和任务分组中的漏洞数量 Badge。
- 撤销全局字体增大 10% 的样式,并删除会话列表中的 Worker 选择框、模型图标和行级 Token 统计。
### Agent
#### 新增的功能
- 新任务可直接关联多个已有任务,实时只读继承直接来源任务的目标、事实、漏洞、已完成意图、资产范围和黑板上下文。
- 黑板读取工具支持按需查询来源任务节点、事实、漏洞和执行轨迹;继承节点带来源标记且所有写工具拒绝修改。
- 企业范围工具支持域名、URL、IP、CIDR、ICP 和企业关键词;关键词仅作为 Agent 范围提示,不参与资产自动归属。
- 漏洞深入操作会在原任务创建带资产锚点和 `derived_from` 边的高优先级人工 Worker 意图。
#### 修改的功能
- 来源任务意图不进入新任务 frontier,新任务继续使用独立 exploration、执行队列、工作目录和对话历史。
- 主 Agent 对话改为服务端活动流驱动,发送失败可恢复输入;打开会话时自动贴底并在详情懒加载后保持最后回复可见。
- 任务暂停会终止当前 Main Agent、Planner 和 Worker 调用,但不阻止暂停期间用户主动发起新的主 Agent 编排会话。
- 取消、关闭和流式中断统一保留真实终止原因、已生成内容、运行轮数、耗时、Token 与未返回工具调用。
#### 删除的功能
- 删除主 Agent 客户端乐观消息回显,避免暂停、失败或并发活动下出现重复消息和跨会话串流。
## [0.3.2] - 2026-08-20
### Added
@@ -43,5 +137,6 @@
- [@Autumn-27](https://github.com/Autumn-27)
- [@neouks](https://github.com/neouks)
[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.2...HEAD
[Unreleased]: https://github.com/Autumn-27/ARTEX/compare/v0.3.3...HEAD
[0.3.3]: https://github.com/Autumn-27/ARTEX/compare/v0.3.2...v0.3.3
[0.3.2]: https://github.com/Autumn-27/ARTEX/compare/v0.3.1...v0.3.2
+16
View File
@@ -108,6 +108,22 @@ CGO_ENABLED=0 go build -tags embedui -o artex ./cmd/artex
./artex
```
### 方式五:构建跨平台 Release 压缩包
`build.sh` 会先构建并嵌入前端,再使用 Go linker 去除调试信息,并将发布文件压缩为 zip。Release 模式默认生成 Linux amd64/arm64、macOS amd64/arm64 和 Windows amd64 的 zip 包:
```bash
./build.sh --release
# 产物:dist/artex-0.3.3-*.zip
```
UPX 自解压二进制可能与部分 Linux 内核、虚拟化环境或安全策略不兼容,因此默认不启用。可用 `ARTEX_TARGETS` 自定义目标;确认目标运行环境兼容时,可显式传入 `--upx` 进一步缩小二进制:
```bash
ARTEX_TARGETS=linux/amd64,windows/amd64 ./build.sh --release
./build.sh --target linux/amd64 --upx
```
---
## 更新升级
+77 -1
View File
@@ -5,6 +5,7 @@ import (
"encoding/json"
"fmt"
"testing"
"time"
"github.com/Autumn-27/artex/db"
actool "github.com/Autumn-27/norma/tool"
@@ -23,6 +24,73 @@ func callReadJSON(t *testing.T, tool actool.CoreTool, input string) any {
return out
}
func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
d := testDB(t)
defer d.Close()
companies := d.Companies()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("overview-scope-%d", time.Now().UnixNano()), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = companies.DeleteCompany(companyID) })
domain := fmt.Sprintf("overview-scope-%d.invalid", companyID)
ip := fmt.Sprintf("2001:db8:%x::42", companyID%0xffff)
cidr := fmt.Sprintf("2001:db8:%x:1::/64", companyID%0xffff)
icp := fmt.Sprintf("京 ICP 备 %d 号", companyID)
keyword := fmt.Sprintf("Scope Company %d", companyID)
inputs := []db.ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "ip", Value: ip},
{Kind: "cidr", Value: cidr},
{Kind: "icp", Value: icp},
{Kind: "keyword", Value: keyword},
}
added, skipped, invalid, scopeErrors := companies.AddScopeInputs(companyID, inputs, "task context test")
if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 {
t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors)
}
task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
tools := NewToolSet(d.Exploration(task.ExplorationID), "planner")
tools.SetTaskID(task.ID)
tools.SetAssetStore(d.Assets(), companies)
overview := tools.graphOverviewData()
coverage, ok := overview["coverage"].(map[string]any)
if !ok {
t.Fatalf("coverage missing: %#v", overview["coverage"])
}
scopeRows, ok := coverage["scope"].([]map[string]any)
if !ok || len(scopeRows) != 1 {
t.Fatalf("task scope missing: %#v", coverage["scope"])
}
companyScope, ok := scopeRows[0]["company_scope"].([]map[string]any)
if !ok || len(companyScope) != len(inputs) {
t.Fatalf("company scope not expanded: %#v", scopeRows[0])
}
kinds := make(map[string]string, len(companyScope))
for _, rule := range companyScope {
kinds[fmt.Sprint(rule["kind"])] = fmt.Sprint(rule["value"])
}
for _, input := range inputs {
if kinds[input.Kind] != input.Value {
t.Errorf("scope %s=%q want %q", input.Kind, kinds[input.Kind], input.Value)
}
}
keywords, ok := scopeRows[0]["company_keywords"].([]string)
if !ok || len(keywords) != 1 || keywords[0] != keyword {
t.Fatalf("company keywords missing: %#v", scopeRows[0]["company_keywords"])
}
}
func TestBlackboardToolsReadDirectSources(t *testing.T) {
d := testDB(t)
defer d.Close()
@@ -125,7 +193,15 @@ func TestBlackboardToolsReadDirectSources(t *testing.T) {
if related[0]["source_task_id"] != source.ID || related[0]["inherited"] != true {
t.Fatalf("related source provenance: %#v", related[0])
}
if facts, ok := related[0]["recent_facts"].([]map[string]any); !ok || len(facts) == 0 || facts[0]["id"] != sourceFact {
recentFacts, ok := related[0]["recent_facts"].([]map[string]any)
foundSourceFact := false
for _, fact := range recentFacts {
if fact["id"] == sourceFact {
foundSourceFact = true
break
}
}
if !ok || !foundSourceFact {
t.Fatalf("related fact summary: %#v", related[0]["recent_facts"])
}
if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding {
+16 -2
View File
@@ -66,6 +66,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
tbuf.WriteString(text)
}
lastTool := &runTrace{startedAt: time.Now()}
var lastUsage *llm.Usage
var finalText string
var rerr error
@@ -74,10 +75,10 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
flush()
if ctx.Err() != nil { // engine/user cancellation, not a provider failure
sum, detail := terminalText(ctx, &harness.Terminal{Reason: reason, Err: ctx.Err()}, lastTool)
rec(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail})
rec(activityWithUsage(db.Activity{Kind: "result", Summary: firstLine(sum, 400), Detail: detail}, lastUsage))
return finalText, reason, ctx.Err()
}
rec(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()})
rec(activityWithUsage(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()}, lastUsage))
return finalText, reason, err
}
switch ev.Kind {
@@ -111,6 +112,7 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
// buffered final-answer text must stay for the KindResult de-dup.
if ev.Usage != nil {
u := *ev.Usage
lastUsage = &u
rec(db.Activity{Kind: "usage",
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
@@ -145,6 +147,18 @@ func captureRunSession(ctx context.Context, s *agentcore.Session, input string,
return finalText, reason, rerr
}
func activityWithUsage(activity db.Activity, usage *llm.Usage) db.Activity {
if usage == nil {
return activity
}
u := *usage
activity.InputTokens = &u.InputTokens
activity.OutputTokens = &u.OutputTokens
activity.CacheReadTokens = &u.CacheReadTokens
activity.CacheWriteTokens = &u.CacheWriteTokens
return activity
}
// blocksText concatenates the text of a tool-result's content blocks.
func blocksText(blocks []llm.ContentBlock) string {
var b strings.Builder
+87
View File
@@ -0,0 +1,87 @@
package agent
import (
"context"
"errors"
"iter"
"testing"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/agentcore"
"github.com/Autumn-27/norma/llm"
)
type captureUsageProvider struct {
stream func(context.Context, func(llm.StreamEvent, error) bool)
}
func (p captureUsageProvider) Stream(ctx context.Context, _ llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) { p.stream(ctx, yield) }
}
func TestCaptureRunPersistsUsageOnProviderFailure(t *testing.T) {
wantErr := errors.New("provider failed after reporting usage")
provider := captureUsageProvider{stream: func(_ context.Context, yield func(llm.StreamEvent, error) bool) {
yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 11, CacheReadTokens: 3}}, nil)
yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "partial"}, nil)
yield(llm.StreamEvent{Type: llm.SEMessageDelta, Usage: llm.Usage{OutputTokens: 7, CacheWriteTokens: 2}}, nil)
yield(llm.StreamEvent{}, wantErr)
}}
var activities []db.Activity
_, _, err := captureRun(context.Background(), agentcore.Options{Provider: provider, MaxTurns: 1}, "test", func(a db.Activity) {
activities = append(activities, a)
})
if !errors.Is(err, wantErr) {
t.Fatalf("captureRun error=%v, want %v", err, wantErr)
}
assertCapturedResultUsage(t, activities, 11, 7, 3, 2)
}
func TestCaptureRunPersistsUsageOnCancellation(t *testing.T) {
started := make(chan struct{})
provider := captureUsageProvider{stream: func(ctx context.Context, yield func(llm.StreamEvent, error) bool) {
yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 13, CacheReadTokens: 5}}, nil)
yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "partial"}, nil)
yield(llm.StreamEvent{Type: llm.SEMessageDelta, Usage: llm.Usage{OutputTokens: 9, CacheWriteTokens: 4}}, nil)
close(started)
<-ctx.Done()
yield(llm.StreamEvent{}, ctx.Err())
}}
ctx, cancel := context.WithCancel(context.Background())
var activities []db.Activity
done := make(chan error, 1)
go func() {
_, _, err := captureRun(ctx, agentcore.Options{Provider: provider, MaxTurns: 1}, "test", func(a db.Activity) {
activities = append(activities, a)
})
done <- err
}()
<-started
cancel()
if err := <-done; !errors.Is(err, context.Canceled) {
t.Fatalf("captureRun error=%v, want context canceled", err)
}
assertCapturedResultUsage(t, activities, 13, 9, 5, 4)
}
func assertCapturedResultUsage(t *testing.T, activities []db.Activity, input, output, read, write int) {
t.Helper()
results := 0
for _, activity := range activities {
if activity.Kind != "result" {
continue
}
results++
if activity.InputTokens == nil || *activity.InputTokens != input ||
activity.OutputTokens == nil || *activity.OutputTokens != output ||
activity.CacheReadTokens == nil || *activity.CacheReadTokens != read ||
activity.CacheWriteTokens == nil || *activity.CacheWriteTokens != write {
t.Fatalf("result usage=%+v, want input=%d output=%d read=%d write=%d", activity, input, output, read, write)
}
}
if results != 1 {
t.Fatalf("result activity count=%d, activities=%+v", results, activities)
}
}
+36
View File
@@ -392,6 +392,42 @@ func (t *ToolSet) graphOverviewData() map[string]any {
e["value"] = r.Net
case r.CompanyID != nil:
e["company_id"] = *r.CompanyID
if t.cs != nil {
if company, err := t.cs.GetCompany(*r.CompanyID); err == nil && company != nil {
e["company_name"] = company.Name
}
if rules, err := t.cs.GetScope(*r.CompanyID); err == nil {
keywords := make([]string, 0)
companyScope := make([]map[string]any, 0, len(rules))
for _, rule := range rules {
value := rule.Raw
if value == "" {
switch rule.Kind {
case "domain":
value = rule.Domain
case "ip", "cidr":
value = rule.Net
default:
value = rule.Value
}
}
entry := map[string]any{"kind": rule.Kind, "value": value}
if rule.Reason != "" {
entry["reason"] = rule.Reason
}
companyScope = append(companyScope, entry)
if rule.Kind == "keyword" && rule.Raw != "" {
keywords = append(keywords, rule.Raw)
}
}
if len(companyScope) > 0 {
e["company_scope"] = companyScope
}
if len(keywords) > 0 {
e["company_keywords"] = keywords
}
}
}
}
scope = append(scope, e)
}
+4 -4
View File
@@ -293,14 +293,14 @@ func (t *ToolSet) insertAssets() actool.CoreTool {
func (t *ToolSet) addCompanyScope() actool.CoreTool {
return writeTool(
"add_company_scope",
"把根域名/IP/IP段(CIDR)加入某公司的【资产范围】——加入后系统自动认领所有命中的资产(现有+未来)。\n"+
"把域名/IP/CIDR/ICP备案/企业关键词加入某公司的【资产范围】——域名、网络和ICP会自动认领命中的资产,关键词只提供给Agent作为范围提示。\n"+
"公司名唯一:company 不存在则新建,已存在则复用(只把范围并进去)。\n"+
"scope 一行一条,支持:根域名 / 单个 IP / CIDR 网段(如 1.2.3.0/24)。\n"+
"scope 一行一条,系统自动识别:根域名 / URL / 单个 IP / CIDR 网段 / ICP备案 / 企业关键词。\n"+
"务必给 reason 说明归属依据(whois/证书/ASN 等)。\n"+
"护栏:拒绝裸 TLD 与过宽网段(IPv4<=/16、IPv6<=/32),非法行会被跳过并在 errors 返回。",
"护栏:拒绝裸 TLD 与过宽网段(IPv4前缀需为/16-/32、IPv6前缀需为/32-/128),非法行会被跳过并在 errors 返回。",
obj(map[string]any{
"company": str("公司名(不存在则新建、存在则复用;名称唯一)"),
"scope": str("资产范围,一行一条:根域名 / IP / CIDR"),
"scope": str("资产范围,一行一条:域名 / URL / IP / CIDR / ICP备案 / 企业关键词"),
"reason": str("归属依据(证据/来源),务必填写"),
"logo": str("公司图标 URL(可选;仅新建公司时生效)"),
}, "company", "scope"),
+230 -46
View File
@@ -1,65 +1,249 @@
#!/usr/bin/env bash
# ARTEX 单二进制构建脚本:静态导出前端并嵌入 Go 二进制。
# ARTEX cross-platform release builder.
#
# 可选环境变量:
# ARTEX_TARGET_OS=linux 目标系统,默认当前系统
# ARTEX_TARGET_ARCH=amd64 目标架构,默认当前架构
# ARTEX_BUILD_VERSION=v0.3.0 构建版本,默认取 git describe
# ARTEX_OUTPUT=/path/to/artex 输出路径
# ARTEX_SKIP_NPM_CI=1 跳过 npm ci,复用现有依赖
# ARTEX_GOSUMDB=sum.golang.org Go checksum database
# The default mode builds one target and embeds the already-exported frontend.
# `./build.sh --release` builds and packages all supported desktop/server targets.
#
# Environment variables:
# ARTEX_TARGET_OS=linux One target OS in single-target mode.
# ARTEX_TARGET_ARCH=amd64 One target arch in single-target mode.
# ARTEX_TARGETS=linux/amd64,... Comma-separated targets for multi-target mode.
# ARTEX_BUILD_VERSION=0.3.3 Version embedded in the binary and archive name.
# ARTEX_OUTPUT=/path/to/artex Explicit binary path in single-target mode.
# ARTEX_OUTPUT_DIR=dist Directory for default binary paths.
# ARTEX_PACKAGE=1 Create a zip archive for each target.
# ARTEX_PACKAGE_DIR=dist Directory for release archives.
# ARTEX_COMPRESS=off UPX mode: off, auto, or required.
# ARTEX_UPX_ARGS="--best --lzma" Arguments passed to UPX.
# ARTEX_SKIP_FRONTEND=1 Reuse server/webui/dist (for CI artifact builds).
# ARTEX_SKIP_NPM_CI=1 Skip npm ci while rebuilding the frontend.
# ARTEX_GOSUMDB=sum.golang.org Go checksum database.
set -euo pipefail
cd "$(cd "$(dirname "$0")" && pwd)"
info() { printf '\033[36m[*]\033[0m %s\n' "$*"; }
ok() { printf '\033[32m[+]\033[0m %s\n' "$*"; }
warn() { printf '\033[33m[!]\033[0m %s\n' "$*" >&2; }
die() { printf '\033[31m[x]\033[0m %s\n' "$*" >&2; exit 1; }
usage() {
cat <<'EOF'
用法:
./build.sh 编译当前系统当前架构
./build.sh --target linux/amd64 编译一个指定目标
./build.sh --release 编译并打包全部支持的目标
选项:
--release 构建 Linux、macOS、Windows 的 amd64/arm64 目标并生成 zip
--target OS/ARCH 设置单个目标,例如 windows/amd64
--upx 强制使用 UPX 压缩二进制(可能影响部分 Linux 环境兼容性)
--no-compress 不使用 UPX,仅使用 Go linker 裁剪并压缩 zip
--help 显示帮助
多目标列表可通过 ARTEX_TARGETS 覆盖,例如:
ARTEX_TARGETS=linux/amd64,windows/amd64 ./build.sh --release
EOF
}
RELEASE_TARGETS_DEFAULT="linux/amd64,linux/arm64,darwin/amd64,darwin/arm64,windows/amd64"
ARTEX_RELEASE="${ARTEX_RELEASE:-0}"
ARTEX_COMPRESS="${ARTEX_COMPRESS:-off}"
ARTEX_PACKAGE="${ARTEX_PACKAGE:-0}"
while [ "$#" -gt 0 ]; do
case "$1" in
--release)
ARTEX_RELEASE=1
ARTEX_PACKAGE=1
shift
;;
--target)
[ "$#" -ge 2 ] || die "--target 需要 OS/ARCH 参数"
target_arg="$2"
case "$target_arg" in
*/*)
ARTEX_TARGET_OS="${target_arg%%/*}"
ARTEX_TARGET_ARCH="${target_arg##*/}"
ARTEX_TARGETS="$target_arg"
;;
*) die "目标必须是 OS/ARCH,例如 linux/amd64" ;;
esac
shift 2
;;
--no-compress)
ARTEX_COMPRESS=0
shift
;;
--upx)
ARTEX_COMPRESS=required
shift
;;
--help|-h)
usage
exit 0
;;
*) die "未知参数:$1(使用 --help 查看用法)" ;;
esac
done
command -v go >/dev/null 2>&1 || die "未检测到 Go(项目需要 Go 1.26 或更高版本)"
command -v npm >/dev/null 2>&1 || die "未检测到 npm(前端静态构建需要 Node.js/npm)"
command -v rsync >/dev/null 2>&1 || die "未检测到 rsync"
ARTEX_GOSUMDB="${ARTEX_GOSUMDB:-sum.golang.org}"
ARTEX_TARGET_OS="${ARTEX_TARGET_OS:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOOS)}"
ARTEX_TARGET_ARCH="${ARTEX_TARGET_ARCH:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOARCH)}"
ARTEX_BUILD_VERSION="${ARTEX_BUILD_VERSION:-0.3.3}"
# Release tags are commonly passed as v0.3.3; keep the binary version consistent.
ARTEX_BUILD_VERSION="${ARTEX_BUILD_VERSION#v}"
ARTEX_OUTPUT_DIR="${ARTEX_OUTPUT_DIR:-dist}"
ARTEX_PACKAGE_DIR="${ARTEX_PACKAGE_DIR:-$ARTEX_OUTPUT_DIR}"
ARTEX_UPX_ARGS="${ARTEX_UPX_ARGS:---best --lzma}"
if [ -z "${ARTEX_BUILD_VERSION:-}" ]; then
if command -v git >/dev/null 2>&1 && git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
ARTEX_BUILD_VERSION="$(git describe --tags --always --dirty)"
else
ARTEX_BUILD_VERSION="dev"
if [ "${ARTEX_RELEASE}" = "1" ]; then
ARTEX_TARGETS="${ARTEX_TARGETS:-$RELEASE_TARGETS_DEFAULT}"
else
ARTEX_TARGET_OS="${ARTEX_TARGET_OS:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOOS)}"
ARTEX_TARGET_ARCH="${ARTEX_TARGET_ARCH:-$(GOSUMDB="$ARTEX_GOSUMDB" go env GOARCH)}"
ARTEX_TARGETS="${ARTEX_TARGETS:-${ARTEX_TARGET_OS}/${ARTEX_TARGET_ARCH}}"
fi
if [ "${ARTEX_SKIP_FRONTEND:-0}" = "1" ]; then
[ -d server/webui/dist ] || die "ARTEX_SKIP_FRONTEND=1 但 server/webui/dist 不存在"
else
command -v npm >/dev/null 2>&1 || die "未检测到 npm(前端静态构建需要 Node.js/npm)"
command -v rsync >/dev/null 2>&1 || die "未检测到 rsync"
info "构建前端静态资源"
if [ "${ARTEX_SKIP_NPM_CI:-0}" != "1" ]; then
(cd web && npm ci)
fi
(cd web && npm run build:static)
info "同步前端资源到 server/webui/dist"
mkdir -p server/webui/dist
rsync -a --delete web/out/ server/webui/dist/
fi
ARTEX_BINARY_NAME="artex"
[ "$ARTEX_TARGET_OS" = "windows" ] && ARTEX_BINARY_NAME="artex.exe"
ARTEX_OUTPUT="${ARTEX_OUTPUT:-dist/artex-${ARTEX_TARGET_OS}-${ARTEX_TARGET_ARCH}/${ARTEX_BINARY_NAME}}"
compress_binary() {
binary="$1"
goos="$2"
case "$ARTEX_COMPRESS" in
0|off|false|none)
info "跳过 UPX:$binary"
return 0
;;
auto|required|true|1) ;;
*) die "ARTEX_COMPRESS 必须是 off、auto 或 required" ;;
esac
info "构建前端静态资源"
if [ "${ARTEX_SKIP_NPM_CI:-0}" != "1" ]; then
(cd web && npm ci)
fi
(cd web && npm run build:static)
info "同步前端资源到 server/webui/dist"
mkdir -p server/webui/dist
rsync -a --delete web/out/ server/webui/dist/
info "编译 ${ARTEX_TARGET_OS}/${ARTEX_TARGET_ARCH},版本 ${ARTEX_BUILD_VERSION}"
mkdir -p "$(dirname "$ARTEX_OUTPUT")"
GOSUMDB="$ARTEX_GOSUMDB" \
CGO_ENABLED=0 \
GOOS="$ARTEX_TARGET_OS" \
GOARCH="$ARTEX_TARGET_ARCH" \
go build \
-tags embedui \
-trimpath \
-ldflags "-s -w -X main.version=${ARTEX_BUILD_VERSION}" \
-o "$ARTEX_OUTPUT" \
./cmd/artex
ok "编译完成:$ARTEX_OUTPUT"
if command -v file >/dev/null 2>&1; then
file "$ARTEX_OUTPUT"
if ! command -v upx >/dev/null 2>&1; then
if [ "$ARTEX_COMPRESS" = "required" ]; then
die "ARTEX_COMPRESS=required 但未检测到 upx"
fi
warn "未检测到 upx,保留 linker 压缩结果:$binary"
return 0
fi
before=$(wc -c < "$binary" | tr -d ' ')
upx_args="$ARTEX_UPX_ARGS"
[ "$goos" = "darwin" ] && upx_args="$upx_args --force-macos"
# shellcheck disable=SC2086
if ! upx $upx_args -- "$binary"; then
if [ "$ARTEX_COMPRESS" = "required" ]; then
die "UPX 压缩失败:$binary"
fi
warn "UPX 不支持该目标格式,保留未压缩二进制:$binary"
return 0
fi
after=$(wc -c < "$binary" | tr -d ' ')
ok "UPX 压缩完成:$binary (${before} -> ${after} bytes)"
}
package_binary() {
binary="$1"
goos="$2"
goarch="$3"
package_name="artex-${ARTEX_BUILD_VERSION}-${goos}-${goarch}"
package_root="${ARTEX_PACKAGE_DIR}/${package_name}"
archive="${ARTEX_PACKAGE_DIR}/${package_name}.zip"
command -v zip >/dev/null 2>&1 || die "打包需要 zip"
rm -rf "$package_root" "$archive"
mkdir -p "$package_root"
cp "$binary" "$package_root/"
cp -R skills "$package_root/"
cp config.example.json "$package_root/"
if [ -f README.md ]; then cp README.md "$package_root/"; fi
(cd "$ARTEX_PACKAGE_DIR" && zip -q -r -9 "$(basename "$archive")" "$(basename "$package_root")")
rm -rf "$package_root"
ok "Release 压缩包:$archive"
}
build_target() {
target="$1"
case "$target" in
*/*) ;;
*) die "无效目标:$target(必须是 OS/ARCH)" ;;
esac
goos="${target%%/*}"
goarch="${target##*/}"
case "$goos" in
linux|darwin|windows) ;;
*) die "不支持的系统:$goos(支持 linux、darwin、windows)" ;;
esac
binary_name="artex"
[ "$goos" = "windows" ] && binary_name="artex.exe"
if [ -n "${ARTEX_OUTPUT:-}" ] && [ "$ARTEX_RELEASE" != "1" ]; then
output="$ARTEX_OUTPUT"
else
output="${ARTEX_OUTPUT_DIR}/artex-${goos}-${goarch}/${binary_name}"
fi
mkdir -p "$(dirname "$output")"
info "编译 ${goos}/${goarch},版本 ${ARTEX_BUILD_VERSION}"
GOSUMDB="$ARTEX_GOSUMDB" \
CGO_ENABLED=0 \
GOOS="$goos" \
GOARCH="$goarch" \
go build \
-tags embedui \
-trimpath \
-ldflags "-s -w -buildid= -X main.version=${ARTEX_BUILD_VERSION}" \
-o "$output" \
./cmd/artex
compress_binary "$output" "$goos"
if command -v file >/dev/null 2>&1; then file "$output"; fi
if [ "$ARTEX_PACKAGE" = "1" ]; then package_binary "$output" "$goos" "$goarch"; fi
ok "编译完成:$output"
}
write_checksums() {
[ "$ARTEX_PACKAGE" = "1" ] || return 0
checksum_file="$ARTEX_PACKAGE_DIR/SHA256SUMS"
if command -v sha256sum >/dev/null 2>&1; then
(cd "$ARTEX_PACKAGE_DIR" && for archive in *.zip; do sha256sum "$archive"; done > "$(basename "$checksum_file")")
elif command -v shasum >/dev/null 2>&1; then
(cd "$ARTEX_PACKAGE_DIR" && for archive in *.zip; do shasum -a 256 "$archive"; done > "$(basename "$checksum_file")")
else
warn "未检测到 sha256sum 或 shasum,跳过 SHA256SUMS"
return 0
fi
ok "校验文件:$checksum_file"
}
mkdir -p "$ARTEX_OUTPUT_DIR"
if [ "$ARTEX_PACKAGE" = "1" ]; then mkdir -p "$ARTEX_PACKAGE_DIR"; fi
old_ifs="$IFS"
IFS=','
read -r -a targets <<< "$ARTEX_TARGETS"
IFS="$old_ifs"
[ "${#targets[@]}" -gt 0 ] || die "ARTEX_TARGETS 不能为空"
for target in "${targets[@]}"; do
target="${target//[[:space:]]/}"
[ -n "$target" ] || continue
build_target "$target"
done
if [ "$ARTEX_PACKAGE" = "1" ]; then
write_checksums
info "Release 包已生成于:$ARTEX_PACKAGE_DIR"
fi
+3 -3
View File
@@ -21,9 +21,9 @@ import (
"github.com/Autumn-27/artex/server"
)
// version is the build version, injected at release time via
// -ldflags "-X main.version=<tag>". Defaults to "dev" for local builds.
var version = "dev"
// version is the application version. Release builds may override it via
// -ldflags "-X main.version=<version>".
var version = "0.3.3"
const banner = `
_ ____ _____ _______ __
+16 -1
View File
@@ -457,6 +457,18 @@ func isIdentChar(c byte) bool {
// ── QueryDSL ─────────────────────────────────────────────────────────────────
// ValidateDSL parses and compiles a DSL expression without touching the
// database. HTTP callers use it to distinguish client syntax errors from query
// failures, which must remain server errors.
func ValidateDSL(dsl string) error {
node, err := ParseDSL(dsl)
if err != nil {
return err
}
_, _, err = buildDSLWhere(node)
return err
}
// CountDSL returns the total number of assets matching a DSL expression (and optional
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
@@ -483,6 +495,9 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
node, err := ParseDSL(dsl)
if err != nil {
return nil, err
@@ -497,7 +512,7 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
}
args = append(args, limit, offset)
q := assetSelectCols + " WHERE " + where +
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
fmt.Sprintf(" ORDER BY last_seen DESC, id DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
+195 -45
View File
@@ -83,6 +83,7 @@ type PortService struct {
type AssetStore struct {
db *DB
company *CompanyStore
tx *sql.Tx
}
// Assets returns the asset store.
@@ -93,6 +94,39 @@ func (d *DB) Assets() *AssetStore {
// Companies returns the company store associated with this asset store.
func (s *AssetStore) Companies() *CompanyStore { return s.company }
// withCompanyScopeMutation serializes scope resolution and every asset write
// that consumes its result in one transaction. Nested asset side effects reuse
// the same transaction through the scoped store.
func (s *AssetStore) withCompanyScopeMutation(fn func(*AssetStore) (int64, error)) (int64, error) {
if s.tx != nil {
return fn(s)
}
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
id, err := fn(scoped)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return id, nil
}
func (s *AssetStore) resolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
if s.tx != nil {
return resolveCompanyWithICP(s.tx, rootDomain, ipStr, icp)
}
return s.company.ResolveCompanyWithICP(rootDomain, ipStr, icp)
}
// =====================================================================
// Helpers
// =====================================================================
@@ -219,7 +253,12 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
companyID, err := s.company.ResolveCompany(domain, "")
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertRootDomain(req)
})
}
companyID, err := s.resolveCompanyWithICP(domain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -237,12 +276,20 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, $4::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, company_source, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, 'scope', $4::bigint[])
ON CONFLICT (domain) WHERE type = 'root_domain' DO UPDATE SET
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
extra = assets.extra || EXCLUDED.extra,
last_seen = now()
@@ -267,8 +314,13 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
if req.IP == "" {
return 0, fmt.Errorf("ip is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertIP(req)
})
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany("", req.IP)
companyID, err := s.resolveCompanyWithICP("", req.IP, "")
if err != nil {
return 0, err
}
@@ -292,9 +344,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, company_source, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, 'scope', $6::bigint[])
ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
bound_domains = (SELECT ARRAY(SELECT DISTINCT unnest(assets.bound_domains || EXCLUDED.bound_domains))),
open_ports = (
@@ -306,7 +358,15 @@ ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
)
),
c_segment = COALESCE(assets.c_segment, EXCLUDED.c_segment),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.IP, csegVal, boundDomains, openPortsJSON, companyID, taskIDs).Scan(&id)
@@ -367,14 +427,19 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertSubdomain(req)
})
}
rootDomain, _ := RootDomain(domain)
if rootDomain == "" {
rootDomain = domain
}
// resolve company by root domain
companyID, err := s.company.ResolveCompany(rootDomain, "")
// Resolve by root domain first, then the asset's exact normalized ICP.
companyID, err := s.resolveCompanyWithICP(rootDomain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -409,7 +474,7 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
cseg := calcCSegment(ipStr)
if companyID == nil && ipStr != "" {
companyID, _ = s.company.ResolveCompany("", ipStr)
companyID, _ = s.resolveCompanyWithICP("", ipStr, "")
}
var taskIDs string
@@ -432,14 +497,22 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
recordType := req.RecordType
recordValueArr := marshalStringArray(req.RecordValue)
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, $9::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, company_source, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, 'scope', $9::bigint[])
ON CONFLICT (domain, COALESCE(record_type,'')) WHERE type = 'subdomain' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
record_value = (SELECT ARRAY(SELECT DISTINCT unnest(assets.record_value || EXCLUDED.record_value))),
extra = assets.extra || EXCLUDED.extra,
@@ -459,7 +532,7 @@ type UpsertAppReq struct {
Category string
Description string
ICP string
CompanyID *int64 // explicit override; nil = no auto attribution for apps
CompanyID *int64 // explicit override; nil = exact ICP auto-attribution when available
TaskID int64
}
@@ -468,6 +541,11 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
if req.Name == "" {
return 0, fmt.Errorf("app name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertApp(req)
})
}
var taskIDs string
if req.TaskID > 0 {
@@ -477,6 +555,7 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
var bundleVal, catVal, descVal, icpVal, companyIDVal any
companySource := "scope"
if req.BundleID != "" {
bundleVal = req.BundleID
}
@@ -491,36 +570,65 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
if req.CompanyID != nil {
companyIDVal = *req.CompanyID
companySource = "explicit"
} else {
companyID, err := s.resolveCompanyWithICP("", "", req.ICP)
if err != nil {
return 0, err
}
if companyID != nil {
companyIDVal = *companyID
}
}
var id int64
var err error
if req.BundleID != "" {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7, $8::bigint[])
ON CONFLICT (bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL DO UPDATE SET
app_name = COALESCE(EXCLUDED.app_name, assets.app_name),
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
} else {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6, $7::bigint[])
ON CONFLICT (app_name) WHERE type = 'app' AND bundle_id IS NULL DO UPDATE SET
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
}
return id, err
}
@@ -548,6 +656,11 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
if req.URL == "" {
return 0, fmt.Errorf("url is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertHTTPService(req)
})
}
normURL := normalizeURL(req.URL)
domain, port, serviceName := parseURL(normURL)
rootDomain, _ := RootDomain(domain)
@@ -556,7 +669,7 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -592,16 +705,16 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, service_type, service_name, domain, ip, port, root_domain,
c_segment, favicon_mmh3, technologies, status_code, content_length, page_title,
auth, company_id, task_ids
auth, company_id, company_source, task_ids
)
VALUES (
'service', $1, 'http', $2, $3, $4, $5, $6,
$7::cidr, $8, $9::text[], $10, $11, $12,
$13::jsonb[], $14, $15::bigint[]
$13::jsonb[], $14, 'scope', $15::bigint[]
)
ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
status_code = COALESCE(EXCLUDED.status_code, assets.status_code),
@@ -618,7 +731,15 @@ ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, serviceName, domainVal, ipVal, nullableInt(port), rootDomain,
@@ -675,6 +796,11 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
if req.ServiceName == "" {
return 0, fmt.Errorf("service_name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertOtherService(req)
})
}
// normalize service_name (lowercase) so the (domain,ip,port,service_name)
// dedup key doesn't split "SSH" and "ssh" into separate rows.
serviceName := strings.ToLower(strings.TrimSpace(req.ServiceName))
@@ -689,7 +815,7 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -721,12 +847,12 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, service_type, service_name, domain, ip, port, root_domain,
c_segment, auth, company_id, task_ids
c_segment, auth, company_id, company_source, task_ids
)
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, $9::bigint[])
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, 'scope', $9::bigint[])
ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other' DO UPDATE SET
domain = COALESCE(EXCLUDED.domain, assets.domain),
ip = COALESCE(EXCLUDED.ip, assets.ip),
@@ -738,7 +864,15 @@ ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE typ
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
serviceName, domainVal, ipVal, req.Port, rootDomainVal,
@@ -799,6 +933,11 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
if req.Method == "" {
return 0, fmt.Errorf("method is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertEndpoint(req)
})
}
method := strings.ToUpper(req.Method)
normURL := normalizeURL(req.URL)
domain, port, _ := parseURL(normURL)
@@ -808,7 +947,7 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -840,12 +979,12 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, method, domain, ip, port, root_domain,
params, company_id, task_ids, c_segment
params, company_id, company_source, task_ids, c_segment
)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, $9::bigint[], $10::cidr)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, 'scope', $9::bigint[], $10::cidr)
ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
params = (
SELECT ARRAY(
@@ -857,7 +996,15 @@ ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, method, domainVal, ipVal, nullableInt(port), rootDomainVal,
@@ -884,6 +1031,9 @@ func (s *AssetStore) QueryByType(typ string, limit, offset int) ([]*Asset, error
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
rows, err := s.db.Query(`
SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
@@ -897,7 +1047,7 @@ SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(method,''), array_to_json(params)::text, extra, last_seen::text
FROM assets
WHERE type = $1
ORDER BY last_seen DESC
ORDER BY last_seen DESC, id DESC
LIMIT $2 OFFSET $3`, typ, limit, offset)
if err != nil {
return nil, err
@@ -954,7 +1104,7 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) {
}
func pageClause(args *[]any, limit, offset int) string {
q := ` ORDER BY last_seen DESC`
q := ` ORDER BY last_seen DESC, id DESC`
if limit > 0 {
*args = append(*args, limit)
q += fmt.Sprintf(` LIMIT $%d`, len(*args))
+50
View File
@@ -4,6 +4,7 @@ import (
"fmt"
"slices"
"testing"
"time"
)
// testSetup opens a DB and returns both stores. Skips if no PG.
@@ -896,3 +897,52 @@ func TestNormalizeURL(t *testing.T) {
}
}
}
func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
d, assets, _ := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
taskID := stamp
marker := fmt.Sprintf("stable-page-%d", stamp)
sharedSeen := time.Date(2026, time.August, 21, 9, 0, 0, 0, time.UTC)
ids := make([]int64, 0, 5)
for i := 0; i < 5; i++ {
var id int64
domain := fmt.Sprintf("%s-%d.invalid", marker, i)
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,task_ids,last_seen)
VALUES ('root_domain',$1,$1,ARRAY[$2]::bigint[],$3) RETURNING id`,
domain, taskID, sharedSeen).Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
defer func() { _, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[])`, ids) }()
want := slices.Clone(ids)
slices.Reverse(want)
collect := func(query func(limit, offset int) ([]*Asset, error)) []int64 {
t.Helper()
var got []int64
for offset := 0; offset < len(ids); offset += 2 {
page, err := query(2, offset)
if err != nil {
t.Fatal(err)
}
for _, asset := range page {
got = append(got, asset.ID)
}
}
return got
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryByTask(taskID, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("task pagination order=%v want=%v", got, want)
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryDSL(marker, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("DSL pagination order=%v want=%v", got, want)
}
}
+462 -109
View File
@@ -2,9 +2,11 @@ package db
import (
"database/sql"
"errors"
"fmt"
"net"
"strings"
"unicode/utf8"
)
// =====================================================================
@@ -25,7 +27,7 @@ type Company struct {
type CompanyWithScope struct {
Company
Scope []ScopeRule `json:"scope"`
AssetCount int `json:"asset_count"`
AssetCount int `json:"asset_count"`
}
// ScopeRule is one company_scope row.
@@ -35,6 +37,7 @@ type ScopeRule struct {
Kind string `json:"kind"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Value string `json:"value,omitempty"`
Raw string `json:"raw"`
Reason string `json:"reason,omitempty"`
}
@@ -42,6 +45,51 @@ type ScopeRule struct {
// CompanyStore operates on the companies + company_scope tables.
type CompanyStore struct{ db *DB }
var (
ErrCompanyNameConflict = errors.New("company name already exists")
ErrCompanyNotFound = errors.New("company not found")
)
const (
// MaxCompanyScopeRules bounds both request processing and the total number
// of persisted rules for one company.
MaxCompanyScopeRules = 256
// Raw and normalized textual scope payloads are bounded by Unicode rune
// count so multi-byte input is treated consistently by the API and DB layer.
MaxCompanyScopeRawRunes = 1024
MaxCompanyScopeValueRunes = 1024
)
// CompanyScopeValidationError identifies a client-correctable scope error.
// Storage and transaction failures are returned as ordinary errors instead.
type CompanyScopeValidationError struct{ Message string }
func (e *CompanyScopeValidationError) Error() string { return e.Message }
// ValidateCompanyScopeInputBounds applies request-wide limits before parsing.
// Store methods call it again so non-HTTP callers cannot bypass the limits.
func ValidateCompanyScopeInputBounds(inputs []ScopeInput) error {
if len(inputs) > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 最多 %d 条", MaxCompanyScopeRules,
)}
}
for i, input := range inputs {
if utf8.RuneCountInString(input.Value) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
}
return nil
}
// Scope writes rebuild derived asset ownership globally, so serialize them to
// ensure the committed attribution always reflects the latest committed rules.
// This key is reserved for company mutations; 7337741001 is the schema lock and
// 7337741002 is the cross-package test-suite lock.
const companyScopeMutationLock int64 = 7337741003
// Companies returns the company store.
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
@@ -69,6 +117,60 @@ RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
return
}
// CreateCompanyWithScope creates a company without updating an existing row.
// The company, its valid initial scope rules, and derived asset attribution are
// committed atomically. Invalid inputs retain the legacy partial-validation
// contract and are reported without preventing valid rules from being stored.
func (s *CompanyStore) CreateCompanyWithScope(name, logo string, inputs []ScopeInput, reason string) (
id int64, added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, 0, nil, err
}
rules, invalid, validationErrors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
nkey := companyNKey(name)
var logoVal any
if logo != "" {
logoVal = logo
}
if err := tx.QueryRow(`
INSERT INTO companies(name, nkey, logo)
VALUES ($1, $2, $3)
ON CONFLICT (nkey) DO NOTHING
RETURNING id`, name, nkey, logoVal).Scan(&id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return 0, 0, 0, invalid, validationErrors, ErrCompanyNameConflict
}
return 0, 0, 0, invalid, validationErrors, err
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, id, rules, reason)
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
}
if err := tx.Commit(); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
return id, added, skipped, invalid, validationErrors, nil
}
// GetCompany returns one company by id (nil if not found).
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
c := &Company{}
@@ -102,12 +204,55 @@ func (s *CompanyStore) UpsertByName(name string) (int64, error) {
return id, err
}
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
// DeleteCompany deletes a company and re-evaluates automatic ownership against
// the remaining companies in the same transaction. Explicitly-owned assets are
// detached by the FK and may then fall back to a remaining scope match.
func (s *CompanyStore) DeleteCompany(id int64) error {
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
_, err := s.DeleteCompanyWithAssets(id, false)
return err
}
// DeleteCompanyWithAssets deletes a company and optionally all of its assets in
// one transaction, then re-evaluates ownership against the remaining companies.
func (s *CompanyStore) DeleteCompanyWithAssets(id int64, deleteAssets bool) (assetsDeleted int64, err error) {
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
if deleteAssets {
res, err := tx.Exec(`DELETE FROM assets WHERE company_id = $1`, id)
if err != nil {
return 0, err
}
assetsDeleted, err = res.RowsAffected()
if err != nil {
return 0, err
}
}
res, err := tx.Exec(`DELETE FROM companies WHERE id = $1`, id)
if err != nil {
return 0, err
}
deleted, err := res.RowsAffected()
if err != nil {
return 0, err
}
if deleted == 0 {
return 0, ErrCompanyNotFound
}
if err := recomputeAttributionTx(tx); err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return assetsDeleted, nil
}
// ListCompanies returns all companies with scope and asset count.
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
rows, err := s.db.Query(`
@@ -149,7 +294,7 @@ ORDER BY c.name`)
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
rows, err := s.db.Query(`
SELECT id, company_id, kind,
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
COALESCE(domain,''), COALESCE(net::text,''), COALESCE(value,''), raw, COALESCE(reason,'')
FROM company_scope
WHERE company_id = $1
ORDER BY id`, companyID)
@@ -157,10 +302,10 @@ ORDER BY id`, companyID)
return nil, err
}
defer rows.Close()
var out []ScopeRule
out := make([]ScopeRule, 0)
for rows.Next() {
var r ScopeRule
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Value, &r.Raw, &r.Reason); err != nil {
return nil, err
}
out = append(out, r)
@@ -171,48 +316,172 @@ ORDER BY id`, companyID)
// AddScope parses and inserts scope lines for a company, then reattributes assets.
// Returns counts of added, skipped, and invalid lines.
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
rule, err := ParseScopeLine(line)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
inserted, err := s.insertScopeRule(companyID, rule, reason)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
if inserted {
added++
} else {
skipped++
}
inputs = append(inputs, ScopeInput{Value: line})
}
// attribute newly-covered assets
if added > 0 {
_ = s.attributeByCompany(companyID)
}
return
return s.AddScopeInputs(companyID, inputs, reason)
}
// insertScopeRule inserts a scope rule. Returns (inserted, error):
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
// AddScopeInputs inserts structured scope rules. Empty kinds use the automatic
// CIDR/IP/ICP/domain/keyword classification 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,184 @@ 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 invalid > 0 {
return 0, invalid, errs, &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围包含 %d 条无效规则,未覆盖原有范围", invalid,
)}
}
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 +679,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 +693,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
}
+167
View File
@@ -1,7 +1,11 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
// cleanup helpers to remove test data
@@ -57,6 +61,70 @@ func TestCompanyUpsertAndGet(t *testing.T) {
}
}
func TestCreateCompanyWithScopeRejectsNormalizedDuplicate(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("Strict Company %d", stamp)
id, added, _, _, validationErrors, err := cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("strict-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
if added != 1 || len(validationErrors) != 0 {
t.Fatalf("initial scope: added=%d errors=%v", added, validationErrors)
}
_, _, _, _, _, err = cs.CreateCompanyWithScope(
" "+strings.ToUpper(strings.ReplaceAll(name, " ", " "))+" ",
"",
[]ScopeInput{{Kind: "domain", Value: fmt.Sprintf("replacement-%d.example", stamp)}},
"test",
)
if !errors.Is(err, ErrCompanyNameConflict) {
t.Fatalf("duplicate create error=%v want ErrCompanyNameConflict", err)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("strict-%d.example", stamp) {
t.Fatalf("duplicate create changed existing scope: %+v", scope)
}
}
func TestCreateCompanyWithScopeRollsBackOnScopeWriteFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Create %d", time.Now().UnixNano())
_, _, _, _, _, err = cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "test")
if err == nil {
t.Fatal("expected scope database write to fail")
}
company, getErr := cs.GetCompanyByName(name)
if getErr != nil {
t.Fatal(getErr)
}
if company != nil {
defer cleanupCompany(d, company.ID)
t.Fatalf("company row survived failed initial scope transaction: %+v", company)
}
}
func TestCompanyScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -187,6 +255,58 @@ func TestUpdateScope(t *testing.T) {
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
t.Errorf("UpdateScope: expected new-domain.com only, got %+v", scope)
}
// Invalid replacement input must not turn a partial validation response into
// a destructive replacement of the existing rules.
added, invalid, validationErrors, err := cs.UpdateScopeInputsChecked(id, []ScopeInput{
{Kind: "domain", Value: "co.uk"},
}, "invalid replacement")
var validationErr *CompanyScopeValidationError
if added != 0 || invalid != 1 || len(validationErrors) != 1 || !errors.As(err, &validationErr) {
t.Fatalf("invalid replacement: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
scope, err = cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
t.Fatalf("invalid replacement changed existing scope: %+v", scope)
}
}
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) {
@@ -217,6 +337,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 {
+15
View File
@@ -0,0 +1,15 @@
package db
import "testing"
func TestCompanyScopeMutationLockDoesNotReuseInfrastructureLocks(t *testing.T) {
reserved := map[string]int64{
"schema migration": 7337741001,
"cross-package test suite": 7337741002,
}
for name, key := range reserved {
if companyScopeMutationLock == key {
t.Fatalf("company scope mutation lock reuses the %s advisory lock key %d", name, key)
}
}
}
+217 -37
View File
@@ -3,20 +3,215 @@ package db
import (
"fmt"
"net"
"net/url"
"strings"
"unicode"
"golang.org/x/net/idna"
"golang.org/x/net/publicsuffix"
)
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
// ParsedScope is one parsed asset-scope entry. Its kind selects Domain, Net, or Value.
// Used internally by the company scope parsers and 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
// uses the same automatic classification as the single-textarea UI.
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), " "))
}
func looksLikeIPAddress(value string) bool {
value = strings.TrimSpace(value)
if strings.Contains(value, "://") || strings.IndexFunc(value, unicode.IsSpace) >= 0 {
return false
}
if strings.Count(value, ":") >= 2 {
// Require an IPv6-looking prefix. This still catches malformed values such
// as 2001:db8::zz without treating ordinary colon-delimited keywords as IPs.
parts := strings.Split(value, ":")
validSegments := 0
for _, part := range parts {
if part == "" {
if validSegments > 0 || strings.HasPrefix(value, "::") {
return true
}
continue
}
if len(part) > 4 {
return false
}
for _, r := range part {
if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') || (r >= 'A' && r <= 'F')) {
return false
}
}
validSegments++
if validSegments >= 2 {
return true
}
}
return false
}
if !strings.Contains(value, ".") {
return false
}
for _, r := range value {
if r != '.' && (r < '0' || r > '9') {
return false
}
}
return true
}
// ParseScopeInput validates an explicitly typed rule. Legacy callers can omit
// Kind and use the same automatic classification as the single-textarea UI.
func ParseScopeInput(input ScopeInput) (ParsedScope, error) {
kind := strings.ToLower(strings.TrimSpace(input.Kind))
raw := strings.TrimSpace(input.Value)
if kind == "" {
return ParseAutoScopeLine(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)
}
}
// ParseAutoScopeLine classifies one untyped textarea line. Network-looking and
// domain-looking values remain strict so malformed ranges do not silently become
// Agent keywords; all other non-empty text is a keyword.
func ParseAutoScopeLine(line string) (ParsedScope, error) {
raw := strings.TrimSpace(line)
if raw == "" {
return ParsedScope{}, fmt.Errorf("空行")
}
if _, _, err := net.ParseCIDR(raw); err == nil {
return ParseScopeLine(raw)
}
if ip := net.ParseIP(raw); ip != nil {
return ParseScopeLine(raw)
}
if slash := strings.LastIndexByte(raw, '/'); slash > 0 {
address := strings.TrimSpace(raw[:slash])
if net.ParseIP(address) != nil || looksLikeIPAddress(address) {
return ParsedScope{Raw: raw}, fmt.Errorf("无效 CIDR: %s", raw)
}
}
if looksLikeIPAddress(raw) {
return ParsedScope{Raw: raw}, fmt.Errorf("无效 IP: %s", raw)
}
looksLikeDomain := strings.Contains(raw, "://") ||
(strings.Contains(raw, ".") && strings.IndexFunc(raw, unicode.IsSpace) < 0)
if looksLikeDomain {
return ParseScopeLine(raw)
}
// 备案号本身不含点号(如 京ICP备12345678号-1)。带点的文本多半掺了域名或版本号,
// 按 ICP 存下来只会得到一条永远匹配不上任何资产的死规则 —— ICP 归属走的是精确
// 相等比较(见 companies.go 的 kind='icp' 归属查询),所以这类文本归为关键词。
lower := strings.ToLower(raw)
if !strings.ContainsAny(raw, "..。") &&
(strings.Contains(lower, "icp") || strings.Contains(raw, "备案")) {
return ParseScopeInput(ScopeInput{Kind: "icp", Value: raw})
}
return ParseScopeInput(ScopeInput{Kind: "keyword", Value: raw})
}
func scopeHostname(raw string) (string, error) {
candidate := strings.TrimSpace(raw)
if candidate == "" {
return "", fmt.Errorf("主机名为空")
}
if strings.HasPrefix(candidate, "//") {
candidate = "http:" + candidate
} else if !strings.Contains(candidate, "://") {
candidate = "http://" + candidate
}
parsed, err := url.Parse(candidate)
if err != nil || parsed.Host == "" {
if err == nil {
err = fmt.Errorf("缺少主机名")
}
return "", err
}
host := strings.TrimSuffix(strings.TrimSpace(parsed.Hostname()), ".")
if host == "" {
return "", fmt.Errorf("主机名为空")
}
if ip := net.ParseIP(host); ip != nil {
return ip.String(), nil
}
host, err = idna.Lookup.ToASCII(host)
if err != nil {
return "", err
}
host = strings.ToLower(host)
if len(host) > 253 {
return "", fmt.Errorf("域名超过 253 个字符")
}
labels := strings.Split(host, ".")
if len(labels) < 2 {
return "", fmt.Errorf("域名至少需要两个标签")
}
for _, label := range labels {
if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' {
return "", fmt.Errorf("域名标签无效")
}
for _, r := range label {
if !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-') {
return "", fmt.Errorf("域名包含无效字符")
}
}
}
return host, nil
}
// 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.
@@ -26,12 +221,8 @@ func ParseScopeLine(line string) (ParsedScope, error) {
if raw == "" {
return r, fmt.Errorf("空行")
}
s := raw
if i := strings.Index(s, "://"); i >= 0 { // strip scheme if a URL was pasted
s = s[i+3:]
}
// CIDR first — it contains '/', which the path-strip below would remove.
if _, ipnet, err := net.ParseCIDR(s); err == nil {
// CIDR first because URL parsing treats its slash as a path separator.
if _, ipnet, err := net.ParseCIDR(raw); err == nil {
ones, bits := ipnet.Mask.Size()
if bits == 32 && ones < 16 {
return r, fmt.Errorf("网段过宽(IPv4 需 >= /16): %s", raw)
@@ -42,12 +233,8 @@ func ParseScopeLine(line string) (ParsedScope, error) {
r.Kind, r.Net = "cidr", ipnet.String()
return r, nil
}
// Strip any path and :port remnants for host-like inputs.
if i := strings.IndexByte(s, '/'); i >= 0 {
s = s[:i]
}
// Single IP.
if ip := net.ParseIP(strings.TrimSpace(s)); ip != nil {
if ip := net.ParseIP(raw); ip != nil {
r.Kind = "ip"
if ip.To4() != nil {
r.Net = ip.String() + "/32"
@@ -56,37 +243,30 @@ func ParseScopeLine(line string) (ParsedScope, error) {
}
return r, nil
}
if host, _, ok := strings.Cut(s, ":"); ok { // host:port → host
s = host
host, err := scopeHostname(raw)
if err != nil {
return r, fmt.Errorf("无法识别为有效域名/IP/CIDR: %s", raw)
}
if ip := net.ParseIP(host); ip != nil {
r.Kind = "ip"
if ip.To4() != nil {
r.Net = ip.String() + "/32"
} else {
r.Net = ip.String() + "/128"
}
return r, nil
}
if looksLikeIPAddress(host) {
return r, fmt.Errorf("无效 IP: %s", raw)
}
if strings.Contains(raw, "-") && strings.Count(raw, ".") >= 6 {
return r, fmt.Errorf("IP 段请用 CIDR 表示(如 1.2.3.0/24): %s", raw)
}
// Domain (registrable). Reject bare TLDs / public suffixes.
d := DomainKey(s)
if d == "" || !strings.Contains(d, ".") {
return r, fmt.Errorf("无法识别为根域名/IP/CIDR: %s", raw)
}
d := DomainKey(host)
if suf, icann := publicsuffix.PublicSuffix(d); icann && suf == d {
return r, fmt.Errorf("不能用裸 TLD 作为范围: %s", raw)
}
r.Kind, r.Domain = "domain", d
return r, nil
}
// ParseScopeLines parses a whole text block (one entry per line), returning the
// valid rules and the invalid lines (with reasons) so the caller can report both.
func ParseScopeLines(text string) (rules []ParsedScope, invalid []string) {
for _, ln := range strings.Split(text, "\n") {
if strings.TrimSpace(ln) == "" {
continue
}
r, err := ParseScopeLine(ln)
if err != nil {
invalid = append(invalid, err.Error())
continue
}
rules = append(rules, r)
}
return rules, invalid
}
+248
View File
@@ -0,0 +1,248 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
domain := fmt.Sprintf("scope-lock-%d.example", stamp)
oldCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock Old %d", stamp), "")
if err != nil {
t.Fatal(err)
}
newCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock New %d", stamp), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain = $1`, domain)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
})
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(oldCompany, []ScopeInput{
{Kind: "domain", Value: domain},
}, "test"); err != nil || added != 1 || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("seed old scope: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
mutation, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer mutation.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(mutation); err != nil {
t.Fatal(err)
}
if _, err := mutation.Exec(`DELETE FROM company_scope WHERE company_id=$1 AND kind='domain' AND domain=$2`, oldCompany, domain); err != nil {
t.Fatal(err)
}
if _, err := insertScopeRuleTx(mutation, newCompany, ParsedScope{
Kind: "domain", Domain: domain, Raw: domain,
}, "test"); err != nil {
t.Fatal(err)
}
started := make(chan struct{})
type result struct {
id int64
err error
}
resultCh := make(chan result, 1)
go func() {
close(started)
id, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
resultCh <- result{id: id, err: err}
}()
<-started
select {
case got := <-resultCh:
t.Fatalf("asset upsert escaped scope mutation lock: %+v", got)
case <-time.After(100 * time.Millisecond):
}
if err := mutation.Commit(); err != nil {
t.Fatal(err)
}
var got result
select {
case got = <-resultCh:
case <-time.After(3 * time.Second):
t.Fatal("asset upsert did not resume after scope mutation committed")
}
if got.err != nil {
t.Fatal(got.err)
}
var companyID *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, got.id).Scan(&companyID); err != nil {
t.Fatal(err)
}
if companyID == nil || *companyID != newCompany {
t.Fatalf("asset retained stale company: got=%v want=%d", companyID, newCompany)
}
}
func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
first, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break First %d", stamp), "")
if err != nil {
t.Fatal(err)
}
second, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break Second %d", stamp), "")
if err != nil {
t.Fatal(err)
}
want := min(first, second)
domain := fmt.Sprintf("tie-break-%d.example", stamp)
segment := uint64(stamp) & 0xffff
network := fmt.Sprintf("2001:db8:%x:1::/64", segment)
ip := fmt.Sprintf("2001:db8:%x:1::42", segment)
icp := fmt.Sprintf("ICP-TIE-%d", stamp)
appName := fmt.Sprintf("Tie Break App %d", stamp)
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
})
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "cidr", Value: network},
{Kind: "icp", Value: icp},
}
for _, companyID := range []int64{second, first} {
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != len(rules) || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("add company %d rules: added=%d invalid=%d validation=%v err=%v", companyID, added, invalid, validationErrors, err)
}
}
assertResolved := func(label string, got *int64, err error) {
t.Helper()
if err != nil {
t.Fatalf("%s resolve: %v", label, err)
}
if got == nil || *got != want {
t.Fatalf("%s resolve=%v want lowest company id %d", label, got, want)
}
}
got, err := companies.ResolveCompany(domain, "")
assertResolved("domain", got, err)
got, err = companies.ResolveCompany("", ip)
assertResolved("network", got, err)
got, err = companies.ResolveCompanyWithICP("", "", icp)
assertResolved("icp", got, err)
rootID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
if err != nil {
t.Fatal(err)
}
ipID, err := assets.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
appID, err := assets.UpsertApp(UpsertAppReq{Name: appName, ICP: icp})
if err != nil {
t.Fatal(err)
}
assertAssets := func(stage string) {
t.Helper()
for _, id := range []int64{rootID, ipID, appID} {
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, id).Scan(&got); err != nil {
t.Fatal(err)
}
if got == nil || *got != want {
t.Fatalf("%s asset %d company=%v want=%d", stage, id, got, want)
}
}
}
assertAssets("live")
if err := companies.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertAssets("recomputed")
}
func TestCompanyScopeLimitsAndCheckedErrors(t *testing.T) {
boundary := strings.Repeat("界", MaxCompanyScopeRawRunes)
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary}}); err != nil {
t.Fatalf("exact raw rune boundary rejected: %v", err)
}
var validationErr *CompanyScopeValidationError
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary + "界"}}); !errors.As(err, &validationErr) {
t.Fatalf("oversized raw value error=%v want CompanyScopeValidationError", err)
}
tooMany := make([]ScopeInput, MaxCompanyScopeRules+1)
if err := ValidateCompanyScopeInputBounds(tooMany); !errors.As(err, &validationErr) {
t.Fatalf("too many rules error=%v want CompanyScopeValidationError", err)
}
d, _, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Limits %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
rules := make([]ScopeInput, MaxCompanyScopeRules)
for i := range rules {
rules[i] = ScopeInput{Kind: "keyword", Value: fmt.Sprintf("limit-%d-%d", stamp, i)}
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != MaxCompanyScopeRules {
t.Fatalf("fill scope capacity: added=%d err=%v", added, err)
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, []ScopeInput{
{Kind: "keyword", Value: fmt.Sprintf("limit-%d-overflow", stamp)},
}, "test"); added != 0 || !errors.As(err, &validationErr) {
t.Fatalf("scope overflow: added=%d err=%v", added, err)
}
var count int
if err := d.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id=$1`, companyID).Scan(&count); err != nil {
t.Fatal(err)
}
if count != MaxCompanyScopeRules {
t.Fatalf("failed capacity mutation changed scope count=%d want=%d", count, MaxCompanyScopeRules)
}
missingID := companyID + 1_000_000_000
if _, _, _, _, err := companies.AddScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing add error=%v want ErrCompanyNotFound", err)
}
if _, _, _, err := companies.UpdateScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing update error=%v want ErrCompanyNotFound", err)
}
if _, err := companies.DeleteCompanyWithAssets(missingID, true); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing delete error=%v want ErrCompanyNotFound", err)
}
}
func TestCompanyScopeCheckedReturnsSystemErrorSeparately(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
companies := d.Companies()
if err := d.Close(); err != nil {
t.Fatal(err)
}
_, _, _, validationErrors, err := companies.AddScopeInputsChecked(1, nil, "test")
if err == nil {
t.Fatal("closed database did not return a system error")
}
if len(validationErrors) != 0 {
t.Fatalf("system error leaked into validation errors: %v", validationErrors)
}
var validationErr *CompanyScopeValidationError
if errors.As(err, &validationErr) || errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("system error misclassified: %v", err)
}
}
+259 -9
View File
@@ -3,6 +3,7 @@ package db
import (
"fmt"
"testing"
"time"
)
// TestParseScopeLine covers classification + guardrails without a DB.
@@ -15,9 +16,9 @@ func TestParseScopeLine(t *testing.T) {
{"example.com", "domain", false},
{"https://sub.example.com/path", "domain", false},
{"1.2.3.4", "ip", false},
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"198.51.100.0/24", "cidr", false},
{"co.uk", "", true}, // bare public suffix
{"co.uk", "", true}, // bare public suffix
{"not a host", "", true},
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
}
@@ -39,6 +40,262 @@ func TestParseScopeLine(t *testing.T) {
}
}
func TestParseAutoScopeLine(t *testing.T) {
cases := []struct {
name string
input string
kind string
normalized string
wantErr bool
}{
{name: "domain", input: "example.com", kind: "domain", normalized: "example.com"},
{name: "url", input: "https://sub.example.com/path", kind: "domain", normalized: "sub.example.com"},
{name: "url query", input: "https://example.com/path?source=x", kind: "domain", normalized: "example.com"},
{name: "url credentials", input: "https://user:[email protected]/path", kind: "domain", normalized: "example.com"},
{name: "url ip", input: "http://203.0.113.10/path", kind: "ip", normalized: "203.0.113.10/32"},
{name: "ipv4", input: "203.0.113.10", kind: "ip", normalized: "203.0.113.10/32"},
{name: "ipv6", input: "2001:db8::10", kind: "ip", normalized: "2001:db8::10/128"},
{name: "cidr", input: "198.51.100.0/24", kind: "cidr", normalized: "198.51.100.0/24"},
{name: "icp latin", input: "京 ICP备 123号", kind: "icp", normalized: "京icp备123号"},
{name: "icp chinese", input: "沪网备案 9988", kind: "icp", normalized: "沪网备案9988"},
{name: "icp domain", input: "icp.example.com", kind: "domain", normalized: "icp.example.com"},
{name: "icp url query", input: "https://example.com/path?icp=1", kind: "domain", normalized: "example.com"},
// 备案号不含点号:掺了域名/版本号的描述性文字归关键词,否则会存成一条
// 永远匹配不上的死 ICP 规则。
{name: "icp with domain text", input: "备案 www.example.com", kind: "keyword", normalized: "备案 www.example.com"},
{name: "icp with version text", input: "某公司 ICP v1.0", kind: "keyword", normalized: "某公司 icp v1.0"},
{name: "icp fullwidth dot", input: "备案 例.com", kind: "keyword", normalized: "备案 例.com"},
{name: "keyword", input: " ACME Security ", kind: "keyword", normalized: "acme security"},
{name: "colon keyword", input: "ACME: Cloud: Security", kind: "keyword", normalized: "acme: cloud: security"},
{name: "empty", input: " ", wantErr: true},
{name: "invalid cidr", input: "10.0.0.0/not-a-prefix", wantErr: true},
{name: "invalid ipv4", input: "999.0.0.1", wantErr: true},
{name: "invalid ipv6", input: "2001:db8::zz", wantErr: true},
{name: "invalid ipv6 cidr", input: "2001:db8::zz/64", wantErr: true},
{name: "overbroad cidr", input: "10.0.0.0/8", wantErr: true},
{name: "bare suffix", input: "co.uk", wantErr: true},
{name: "empty domain label", input: "foo..example.com", wantErr: true},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
rule, err := ParseAutoScopeLine(tc.input)
if tc.wantErr {
if err == nil {
t.Fatalf("ParseAutoScopeLine(%q) = %+v, want error", tc.input, rule)
}
return
}
if err != nil {
t.Fatalf("ParseAutoScopeLine(%q): %v", tc.input, err)
}
if rule.Kind != tc.kind {
t.Fatalf("kind=%q want %q", rule.Kind, tc.kind)
}
got := rule.Value
if rule.Kind == "domain" {
got = rule.Domain
} else if rule.Kind == "ip" || rule.Kind == "cidr" {
got = rule.Net
}
if got != tc.normalized {
t.Fatalf("normalized=%q want %q", got, tc.normalized)
}
})
}
}
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
@@ -84,13 +341,6 @@ func TestCompanyScopeAttribution(t *testing.T) {
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
}
rules, invalid := ParseScopeLines(root + "\n198.51.100.0/24\nco.uk")
if len(rules) != 2 {
t.Fatalf("want 2 valid rules, got %d (%+v)", len(rules), rules)
}
if len(invalid) != 1 {
t.Fatalf("want 1 invalid line, got %v", invalid)
}
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
mustCid := func(id int64, want int64, label string) {
+47 -11
View File
@@ -50,19 +50,32 @@ GROUP BY c.id`)
// Conversation is one ChatGPT-style chat thread bound to an agent key. It lives
// independent of the pentest exploration graph — see schema.sql §I.
type Conversation struct {
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
Pinned bool `json:"pinned"`
PinnedAt *time.Time `json:"pinned_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at`
// ConversationPatch updates only the fields whose pointers are non-nil.
type ConversationPatch struct {
Title *string
Pinned *bool
}
const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at`
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
var c Conversation
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt)
var pinnedAt sql.NullTime
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt)
if pinnedAt.Valid {
c.Pinned = true
c.PinnedAt = &pinnedAt.Time
}
return c, err
}
@@ -127,7 +140,8 @@ func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error {
// ListConversations returns all threads, most-recently-updated first.
func (d *DB) ListConversations() ([]*Conversation, error) {
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`)
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations
ORDER BY (pinned_at IS NOT NULL) DESC, pinned_at DESC NULLS LAST, updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
@@ -155,9 +169,31 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) {
return &c, nil
}
// RenameConversation sets a thread's title.
// UpdateConversation applies a partial title/pin mutation and returns the updated
// row. Pinning an already-pinned conversation preserves its original pin order.
func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`UPDATE conversations SET
title = CASE WHEN $2::boolean THEN $3 ELSE title END,
pinned_at = CASE
WHEN $4::boolean IS NULL THEN pinned_at
WHEN $4::boolean THEN COALESCE(pinned_at, now())
ELSE NULL
END
WHERE id=$1
RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &c, nil
}
// RenameConversation sets a thread's title. Kept for automatic first-message
// titles and compatibility with existing callers.
func (d *DB) RenameConversation(id int64, title string) error {
_, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title)
_, err := d.UpdateConversation(id, ConversationPatch{Title: &title})
return err
}
+69
View File
@@ -0,0 +1,69 @@
package db
import (
"fmt"
"testing"
"time"
)
func TestConversationPinOrderingAndPatch(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
first, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-first-%d", suffix), nil)
if err != nil {
t.Fatal(err)
}
second, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-second-%d", suffix), nil)
if err != nil {
_ = d.DeleteConversation(first.ID)
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(first.ID)
_ = d.DeleteConversation(second.ID)
})
pinned := true
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || !first.Pinned || first.PinnedAt == nil {
t.Fatalf("pin first = %+v, %v", first, err)
}
firstPinnedAt := *first.PinnedAt
time.Sleep(5 * time.Millisecond)
second, err = d.UpdateConversation(second.ID, ConversationPatch{Pinned: &pinned})
if err != nil || second == nil || !second.Pinned {
t.Fatalf("pin second = %+v, %v", second, err)
}
newTitle := "renamed while pinned"
first, err = d.UpdateConversation(first.ID, ConversationPatch{Title: &newTitle, Pinned: &pinned})
if err != nil || first == nil {
t.Fatal(err)
}
if first.Title != newTitle || first.PinnedAt == nil || !first.PinnedAt.Equal(firstPinnedAt) {
t.Fatalf("repeat pin should preserve pin time: %+v (want %v)", first, firstPinnedAt)
}
conversations, err := d.ListConversations()
if err != nil {
t.Fatal(err)
}
positions := map[int64]int{}
for index, conversation := range conversations {
positions[conversation.ID] = index
}
if positions[second.ID] >= positions[first.ID] {
t.Fatalf("newer pin must sort first: second=%d first=%d", positions[second.ID], positions[first.ID])
}
pinned = false
first, err = d.UpdateConversation(first.ID, ConversationPatch{Pinned: &pinned})
if err != nil || first == nil || first.Pinned || first.PinnedAt != nil {
t.Fatalf("unpin first = %+v, %v", first, err)
}
}
+215 -6
View File
@@ -3,11 +3,17 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"time"
)
// ErrIntentStateConflict marks a failed compare-and-set transition. Callers use
// it to distinguish an expected control race (HTTP 409) from a storage failure.
var ErrIntentStateConflict = errors.New("intent state changed concurrently")
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
@@ -70,6 +76,17 @@ type TokenUsage struct {
CacheWriteTokens int `json:"cache_write_tokens"`
}
// SessionTokenUsage is the authoritative completed-run total for one task UI
// session. Worker sessions are keyed by intent rather than the reusable work#N
// executor name, so retries and reassignment remain attached to the same row.
type SessionTokenUsage struct {
Session string `json:"session"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// DailyTokenBucket is one day's global token aggregate across all tasks.
type DailyTokenBucket struct {
Day string `json:"day"` // "YYYY-MM-DD"
@@ -264,6 +281,21 @@ WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, termin
return err
}
// CompareAndSetIntentState transitions one local intent only when its current
// state still matches expected. It is the state boundary used by pause/resume and
// worker settlement, so a stale API read cannot overwrite a concurrent finish.
func (s *ExplorationStore) CompareAndSetIntentState(id int64, expected, state string) (bool, error) {
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
res, err := s.db.Exec(`UPDATE exploration_nodes
SET state=$1, blocked_reason=NULL, completed_at = CASE WHEN $5 THEN now() ELSE NULL END
WHERE id=$2 AND exploration_id=$3 AND kind='intent' AND state=$4`, state, id, s.expID, expected, terminal)
if err != nil {
return false, err
}
n, err := res.RowsAffected()
return n == 1, err
}
// IntentCleanup summarizes the blackboard records removed when a user cancels
// one worker intent. Global assets and traffic are deliberately not part of this
// operation; exploration anchors disappear only because their owning nodes do.
@@ -295,7 +327,7 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, err
}
if state != "running" && state != "paused" {
return out, fmt.Errorf("intent state %s cannot be cancelled", state)
return out, fmt.Errorf("%w: intent state %s cannot be cancelled", ErrIntentStateConflict, state)
}
if err := tx.QueryRow(`SELECT
@@ -304,14 +336,33 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
FROM exploration_edges e
JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil {
AND n.kind IN ('fact','finding')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, s.expID, id, RelYields).Scan(&out.Facts, &out.Findings); err != nil {
return out, err
}
if _, err := tx.Exec(`DELETE FROM findings f USING exploration_edges e, exploration_nodes n
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.id=e.dst_id AND n.exploration_id=e.exploration_id
AND n.kind='finding' AND f.node_id=n.id`, s.expID, id, RelYields); err != nil {
AND n.kind='finding' AND f.node_id=n.id
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)`, s.expID, id, RelYields); err != nil {
return out, err
}
// Activity rows are part of the blackboard cleanup, but their token usage is
// irreversible metering data. Fold completed runs plus the latest unfinished
// usage frame into detached daily result rows before deleting the conversation.
// The rows are excluded from per-intent sessions while whole-task totals and
// the original UTC reporting dates remain accurate.
tokenBuckets, err := intentTokenRollup(tx, s.expID, id)
if err != nil {
return out, err
}
res, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id)
@@ -319,15 +370,36 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, err
}
out.Activities, _ = res.RowsAffected()
for _, bucket := range tokenBuckets {
metadata, _ := json.Marshal(map[string]any{
"cancelled_intent_id": id,
"token_day": bucket.Day.Format(time.DateOnly),
"token_rollup": true,
})
if _, err := tx.Exec(`INSERT INTO activity(
exploration_id, worker, kind, summary, metadata,
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at)
VALUES ($1,'token-ledger','result',$2,$3,$4,$5,$6,$7,$8)`,
s.expID, fmt.Sprintf("已取消意图 #%d 的 Token 计量", id), metadata,
bucket.Usage.InputTokens, bucket.Usage.OutputTokens,
bucket.Usage.CacheReadTokens, bucket.Usage.CacheWriteTokens, bucket.Day); err != nil {
return out, err
}
}
res, err = tx.Exec(`DELETE FROM exploration_nodes
WHERE exploration_id=$1 AND (
id=$2 OR id IN (
SELECT e.dst_id FROM exploration_edges e
JOIN exploration_nodes n ON n.id=e.dst_id AND n.exploration_id=e.exploration_id
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')
)
WHERE e.exploration_id=$1 AND e.src_id=$2 AND e.rel=$3
AND n.kind IN ('fact','finding')
AND NOT EXISTS (
SELECT 1 FROM exploration_edges other
WHERE other.exploration_id=e.exploration_id AND other.dst_id=e.dst_id
AND other.src_id<>$2
)
)
)`, s.expID, id, RelYields)
if err != nil {
return out, err
@@ -339,6 +411,108 @@ func (s *ExplorationStore) CancelIntent(id int64) (IntentCleanup, error) {
return out, tx.Commit()
}
type tokenUsageBucket struct {
Day time.Time
Usage TokenUsage
}
// intentTokenRollup returns exactly-once usage grouped by its original UTC day.
// Each result terminates one run. A result carrying usage is authoritative; an
// older error/cancel result without usage falls back to the latest cumulative
// usage frame since the previous result. A trailing frame represents a run that
// was interrupted before its terminal result could be persisted.
func intentTokenRollup(tx *sql.Tx, explorationID, intentID int64) ([]tokenUsageBucket, error) {
rows, err := tx.Query(`SELECT kind, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at
FROM activity
WHERE exploration_id=$1 AND node_id=$2 AND kind IN ('usage','result')
ORDER BY id`, explorationID, intentID)
if err != nil {
return nil, err
}
defer rows.Close()
byDay := make(map[time.Time]TokenUsage)
add := func(at time.Time, usage TokenUsage) {
if usage.InputTokens == 0 && usage.OutputTokens == 0 &&
usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
return
}
utc := at.UTC()
day := time.Date(utc.Year(), utc.Month(), utc.Day(), 0, 0, 0, 0, time.UTC)
total := byDay[day]
total.add(usage)
byDay[day] = total
}
var pending TokenUsage
var pendingAt time.Time
hasPending := false
for rows.Next() {
var kind string
var input, output, read, write sql.NullInt64
var createdAt time.Time
if err := rows.Scan(&kind, &input, &output, &read, &write, &createdAt); err != nil {
return nil, err
}
current := TokenUsage{
InputTokens: int(input.Int64),
OutputTokens: int(output.Int64),
CacheReadTokens: int(read.Int64),
CacheWriteTokens: int(write.Int64),
}
hasCurrent := input.Valid || output.Valid || read.Valid || write.Valid
if kind == "usage" {
if hasCurrent {
pending, pendingAt, hasPending = current, createdAt, true
}
continue
}
if hasCurrent {
// Preserve a partially populated legacy terminal row by filling only its
// missing dimensions from the latest cumulative usage frame.
if hasPending {
if !input.Valid {
current.InputTokens = pending.InputTokens
}
if !output.Valid {
current.OutputTokens = pending.OutputTokens
}
if !read.Valid {
current.CacheReadTokens = pending.CacheReadTokens
}
if !write.Valid {
current.CacheWriteTokens = pending.CacheWriteTokens
}
}
add(createdAt, current)
} else if hasPending {
add(createdAt, pending)
}
pending, pendingAt, hasPending = TokenUsage{}, time.Time{}, false
}
if err := rows.Err(); err != nil {
return nil, err
}
if hasPending {
add(pendingAt, pending)
}
buckets := make([]tokenUsageBucket, 0, len(byDay))
for day, usage := range byDay {
buckets = append(buckets, tokenUsageBucket{Day: day, Usage: usage})
}
sort.Slice(buckets, func(i, j int) bool { return buckets[i].Day.Before(buckets[j].Day) })
return buckets, nil
}
func (u *TokenUsage) add(other TokenUsage) {
u.InputTokens += other.InputTokens
u.OutputTokens += other.OutputTokens
u.CacheReadTokens += other.CacheReadTokens
u.CacheWriteTokens += other.CacheWriteTokens
}
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), COALESCE(blocked_reason,''), created_at`
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
@@ -761,6 +935,41 @@ func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
return out, rows.Err()
}
// TokenStatsBySession aggregates every persisted completed run, independent of
// activity-history pagination. Main/planner use fixed keys; Worker executions use
// their intent node id, even when a different work#N process handled a retry.
func (s *ExplorationStore) TokenStatsBySession() ([]SessionTokenUsage, error) {
rows, err := s.db.Query(`SELECT
CASE
WHEN COALESCE(a.worker,'')='mainagent' THEN 'main'
WHEN COALESCE(a.worker,'')='planner' THEN 'plan'
WHEN n.id IS NOT NULL THEN 'intent:' || n.id::text
ELSE ''
END AS session_key,
COALESCE(SUM(a.input_tokens),0), COALESCE(SUM(a.output_tokens),0),
COALESCE(SUM(a.cache_read_tokens),0), COALESCE(SUM(a.cache_write_tokens),0)
FROM activity a
LEFT JOIN exploration_nodes n
ON n.id=a.node_id AND n.exploration_id=a.exploration_id AND n.kind='intent'
WHERE a.exploration_id=$1 AND a.kind='result'
AND (COALESCE(a.worker,'') IN ('mainagent','planner') OR n.id IS NOT NULL)
GROUP BY session_key
ORDER BY session_key`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []SessionTokenUsage{}
for rows.Next() {
var u SessionTokenUsage
if err := rows.Scan(&u.Session, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out = append(out, u)
}
return out, rows.Err()
}
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
// Returns items and the new cursor (max id seen).
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
+130
View File
@@ -0,0 +1,130 @@
package db
import (
"fmt"
"testing"
)
func TestTokenStatsBySessionUsesCompleteIntentHistory(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("token sessions", "aggregate complete session history")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "worker session"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
factID, err := store.AddNode(KindFact, map[string]any{"summary": "not a worker"}, 0, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
otherExpID, err := d.CreateExploration("other token sessions", "foreign intent must not match")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, otherExpID)
foreignIntentID, err := d.Exploration(otherExpID).AddIntent(map[string]any{"summary": "foreign worker"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
appendResult := func(nodeID *int64, worker string, input, output, read, write int) {
t.Helper()
if _, appendErr := store.AppendActivity(Activity{
NodeID: nodeID,
Worker: worker,
Kind: "result",
InputTokens: &input,
OutputTokens: &output,
CacheReadTokens: &read,
CacheWriteTokens: &write,
}); appendErr != nil {
t.Fatal(appendErr)
}
}
appendResult(nil, "mainagent", 11, 12, 13, 14)
appendResult(nil, "planner", 21, 22, 23, 24)
// More than ActivityPage's default limit proves this aggregate reads the
// persisted history directly instead of summing the currently loaded page.
const completedRuns = 205
for i := 0; i < completedRuns; i++ {
appendResult(&intentID, fmt.Sprintf("work#%d", i%3+1), 1, 2, 3, 4)
}
// All three rows remain part of the legacy worker/whole-task totals, but none
// represents a local Worker session and therefore none may create intent:*.
appendResult(&factID, "work#fact", 31, 32, 33, 34)
appendResult(nil, "work#missing-node", 41, 42, 43, 44)
appendResult(&foreignIntentID, "work#foreign", 51, 52, 53, 54)
ignoredInput := 1000
if _, err := store.AppendActivity(Activity{
NodeID: &intentID,
Worker: "work#1",
Kind: "usage",
InputTokens: &ignoredInput,
}); err != nil {
t.Fatal(err)
}
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
got := make(map[string]SessionTokenUsage, len(sessions))
for _, usage := range sessions {
got[usage.Session] = usage
}
if len(got) != 3 {
t.Fatalf("sessions = %+v, want only main, plan, and the local intent", sessions)
}
assertSessionTokens(t, got["main"], "main", 11, 12, 13, 14)
assertSessionTokens(t, got["plan"], "plan", 21, 22, 23, 24)
assertSessionTokens(t, got[fmt.Sprintf("intent:%d", intentID)], fmt.Sprintf("intent:%d", intentID),
completedRuns, completedRuns*2, completedRuns*3, completedRuns*4)
workers, err := store.TokenStatsByWorker()
if err != nil {
t.Fatal(err)
}
if len(workers) != 8 { // main, planner, 3 executors, and the 3 deliberately invalid session rows
t.Fatalf("legacy workers changed: got %d entries: %+v", len(workers), workers)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total,
11+21+completedRuns+31+41+51,
12+22+completedRuns*2+32+42+52,
13+23+completedRuns*3+33+43+53,
14+24+completedRuns*4+34+44+54)
}
func assertSessionTokens(t *testing.T, got SessionTokenUsage, session string, input, output, read, write int) {
t.Helper()
if got.Session != session || got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("session %q = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
session, got, input, output, read, write)
}
}
func assertTokenUsage(t *testing.T, got TokenUsage, input, output, read, write int) {
t.Helper()
if got.InputTokens != input || got.OutputTokens != output ||
got.CacheReadTokens != read || got.CacheWriteTokens != write {
t.Fatalf("token total = %+v, want input=%d output=%d cache-read=%d cache-write=%d",
got, input, output, read, write)
}
}
+200 -4
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
@@ -147,6 +148,11 @@ type FindingFilter struct {
Sort string // "severity" | "time"
}
// FindingUnassignedTask is the task filter sentinel for findings whose task is
// absent. That includes rows created without a task and rows retained after their
// originating task was deleted (the findings FK is ON DELETE SET NULL).
const FindingUnassignedTask = "__unassigned__"
// where builds the WHERE clause (shared by the page and count queries) plus its
// positional args. Only equality filters, all parameterized.
func (f FindingFilter) where() (string, []any) {
@@ -163,7 +169,9 @@ func (f FindingFilter) where() (string, []any) {
add("status", f.Status)
add("vulnclass", f.VulnClass)
// task_id 是 bigint 列,按整数比较(不能走上面的文本 add);空/非法值忽略。
if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
if f.TaskID == FindingUnassignedTask {
conds = append(conds, "(f.task_id IS NULL OR t.id IS NULL)")
} else if tid, err := strconv.ParseInt(f.TaskID, 10, 64); err == nil && tid > 0 {
args = append(args, tid)
conds = append(conds, fmt.Sprintf("f.task_id = $%d", len(args)))
}
@@ -185,14 +193,19 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
where, args := f.where()
var total int
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f`+where, args...).Scan(&total); err != nil {
if err := d.QueryRow(`SELECT COUNT(*) FROM findings f LEFT JOIN tasks t ON f.task_id=t.id`+where, args...).Scan(&total); err != nil {
return nil, 0, err
}
// Avoid overflowing (page-1)*pageSize for an arbitrarily large page number.
// Once the requested page is beyond the exact count, no data query is needed.
if total == 0 || page > (total-1)/pageSize+1 {
return []*DBFinding{}, total, nil
}
order := "f.created_at DESC"
order := "f.created_at DESC, f.id DESC"
if f.Sort == "severity" {
// critical > high > medium > low > 其它, then newest first.
order = `CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END DESC, f.created_at DESC`
order = `CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END DESC, f.created_at DESC, f.id DESC`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
q := fmt.Sprintf(`
@@ -210,6 +223,189 @@ func (d *DB) ListFindingsPage(f FindingFilter, page, pageSize int) ([]*DBFinding
return out, total, err
}
// FindingGroup is one task-level bucket in the global findings view. TaskID is
// nil for both findings that never had a task and findings retained after task
// deletion; those records intentionally share one "unassigned/deleted" bucket.
type FindingGroup struct {
TaskID *int64 `json:"task_id"`
TaskDescription string `json:"task_description"`
TaskStatus string `json:"task_status"`
Count int `json:"count"`
Critical int `json:"critical"`
High int `json:"high"`
Medium int `json:"medium"`
Low int `json:"low"`
LastFoundAt time.Time `json:"last_found_at"`
}
// ListFindingGroups returns a page of task groups matching the same filters as
// ListFindingsPage. The group count and finding count are independent totals so
// clients can page groups without losing the exact export/selection count.
func (d *DB) ListFindingGroups(f FindingFilter, page, pageSize int) ([]FindingGroup, int, int, error) {
if page <= 0 {
page = 1
}
if pageSize <= 0 {
pageSize = 10
}
where, args := f.where()
grouped := ` FROM findings f LEFT JOIN tasks t ON f.task_id=t.id` + where +
` GROUP BY t.id, t.description, t.status, t.paused, t.queued`
var groupTotal, findingTotal int
countQuery := `SELECT COUNT(*), COALESCE(SUM(finding_count),0) FROM (` +
`SELECT COUNT(*) AS finding_count` + grouped + `) grouped_findings`
if err := d.QueryRow(countQuery, args...).Scan(&groupTotal, &findingTotal); err != nil {
return nil, 0, 0, err
}
if groupTotal == 0 || page > (groupTotal-1)/pageSize+1 {
return []FindingGroup{}, groupTotal, findingTotal, nil
}
order := "MAX(f.created_at) DESC, t.id DESC NULLS LAST"
if f.Sort == "severity" {
order = `MAX(CASE f.severity WHEN 'critical' THEN 4 WHEN 'high' THEN 3 WHEN 'medium' THEN 2 WHEN 'low' THEN 1 ELSE 0 END) DESC, MAX(f.created_at) DESC, t.id DESC NULLS LAST`
}
pageArgs := append(append([]any{}, args...), pageSize, (page-1)*pageSize)
query := fmt.Sprintf(`SELECT t.id, COALESCE(t.description,''), COALESCE(
CASE
WHEN t.status IN ('done','failed','timeout') THEN t.status
WHEN t.queued THEN 'queued'
WHEN t.paused THEN 'paused'
ELSE t.status
END, ''),
COUNT(*),
COUNT(*) FILTER (WHERE f.severity='critical'),
COUNT(*) FILTER (WHERE f.severity='high'),
COUNT(*) FILTER (WHERE f.severity='medium'),
COUNT(*) FILTER (WHERE f.severity='low'),
MAX(f.created_at)%s
ORDER BY %s LIMIT $%d OFFSET $%d`, grouped, order, len(args)+1, len(args)+2)
rows, err := d.Query(query, pageArgs...)
if err != nil {
return nil, 0, 0, err
}
defer rows.Close()
groups := []FindingGroup{}
for rows.Next() {
var group FindingGroup
var taskID sql.NullInt64
if err := rows.Scan(&taskID, &group.TaskDescription, &group.TaskStatus, &group.Count,
&group.Critical, &group.High, &group.Medium, &group.Low, &group.LastFoundAt); err != nil {
return nil, 0, 0, err
}
if taskID.Valid {
id := taskID.Int64
group.TaskID = &id
}
groups = append(groups, group)
}
return groups, groupTotal, findingTotal, rows.Err()
}
// ErrFindingOriginUnavailable means a retained finding no longer has a live
// owning task and finding node from which a follow-up intent can be derived.
var ErrFindingOriginUnavailable = errors.New("finding origin is no longer available")
// AddFindingFollowUpIntent atomically creates a priority-10 human intent from a
// live finding node, copies that finding's asset anchors, records the
// finding --derived_from--> intent lineage edge, and persists its audit activity.
// The returned activity is the committed row and can be broadcast as-is without
// calling AppendActivity again.
func (s *ExplorationStore) AddFindingFollowUpIntent(findingID, findingNodeID int64, description string, audit Activity) (int64, Activity, error) {
tx, err := s.db.Begin()
if err != nil {
return 0, Activity{}, err
}
defer tx.Rollback()
var liveNodeID int64
err = tx.QueryRow(`SELECT n.id
FROM findings f
JOIN tasks t ON t.id=f.task_id
JOIN exploration_nodes n ON n.id=f.node_id AND n.exploration_id=t.exploration_id
WHERE f.id=$1 AND f.node_id=$2 AND t.exploration_id=$3 AND n.kind='finding'
FOR SHARE OF f, t, n`, findingID, findingNodeID, s.expID).Scan(&liveNodeID)
if err == sql.ErrNoRows {
return 0, Activity{}, ErrFindingOriginUnavailable
}
if err != nil {
return 0, Activity{}, err
}
anchors := []int64{}
anchorRows, err := tx.Query(`SELECT asset_id FROM exploration_anchors WHERE node_id=$1 ORDER BY asset_id`, liveNodeID)
if err != nil {
return 0, Activity{}, err
}
for anchorRows.Next() {
var assetID int64
if err := anchorRows.Scan(&assetID); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
anchors = append(anchors, assetID)
}
if err := anchorRows.Err(); err != nil {
anchorRows.Close()
return 0, Activity{}, err
}
if err := anchorRows.Close(); err != nil {
return 0, Activity{}, err
}
payload := map[string]any{
"summary": description,
"source_finding_id": findingID,
"source_finding_node_id": liveNodeID,
}
if len(anchors) > 0 {
payload["asset_ids"] = anchors
}
raw, err := json.Marshal(payload)
if err != nil {
return 0, Activity{}, err
}
var intentID int64
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
VALUES ($1,'intent',$2,10,'open','human') RETURNING id`, s.expID, raw).Scan(&intentID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id)
SELECT $1, asset_id FROM exploration_anchors WHERE node_id=$2
ON CONFLICT DO NOTHING`, intentID, liveNodeID); err != nil {
return 0, Activity{}, err
}
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id)
VALUES ($1,$2,$3,$4)`, s.expID, liveNodeID, RelDerivedFrom, intentID); err != nil {
return 0, Activity{}, err
}
audit.NodeID = &intentID
if summary := strings.TrimSpace(audit.Summary); summary != "" {
audit.Summary = fmt.Sprintf("%s #%d", summary, intentID)
} else {
audit.Summary = ""
}
metadata := audit.Metadata
if len(metadata) == 0 {
metadata = json.RawMessage(`{}`)
}
if err := tx.QueryRow(`
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, metadata, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13,$14)
RETURNING id, created_at`, s.expID, audit.NodeID, utf8Clean(audit.Worker), utf8Clean(audit.Kind), utf8Clean(audit.Tool), utf8Clean(audit.ToolUseID), audit.IsError,
utf8Clean(audit.Summary), utf8Clean(audit.Detail), metadata, audit.InputTokens, audit.OutputTokens, audit.CacheReadTokens, audit.CacheWriteTokens).
Scan(&audit.ID, &audit.CreatedAt); err != nil {
return 0, Activity{}, err
}
audit.Metadata = metadata
if err := tx.Commit(); err != nil {
return 0, Activity{}, err
}
return intentID, audit, nil
}
// ListFindingsForExport returns findings for the 发现 page 导出功能,携带完整
// report 字段、不分页。ids 非空时按这批 finding id 精确导出(勾选导出),忽略
// filter;ids 为空时按 filter 导出(导出当前筛选/全部)。结果按严重等级降序、
+298
View File
@@ -1,8 +1,12 @@
package db
import (
"encoding/json"
"errors"
"fmt"
"slices"
"testing"
"time"
)
// TestDeleteFinding verifies删除漏洞 removes both the findings row and its
@@ -183,3 +187,297 @@ func TestFindingsPageAndStats(t *testing.T) {
t.Fatalf("SetFindingSeverity(-1): want 0,nil got %d,%v", n, err)
}
}
func TestListFindingsPageUsesStableIDTieBreaker(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_stable_page_%d__", time.Now().UnixNano())
defer d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc)
ids := make([]int64, 0, 7)
for i := 0; i < 7; i++ {
id, addErr := d.AddFinding(0, 0, vc, "", SeverityHigh, fmt.Sprintf("finding %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
ids = append(ids, id)
}
sharedCreatedAt := time.Date(2026, time.August, 21, 8, 30, 0, 0, time.UTC)
if _, err := d.Exec(`UPDATE findings SET created_at=$1 WHERE vulnclass=$2`, sharedCreatedAt, vc); err != nil {
t.Fatal(err)
}
want := make([]int64, len(ids))
for i := range ids {
want[i] = ids[len(ids)-1-i]
}
for _, sort := range []string{"time", "severity"} {
t.Run(sort, func(t *testing.T) {
var got []int64
for page := 1; page <= 3; page++ {
items, total, pageErr := d.ListFindingsPage(FindingFilter{VulnClass: vc, Sort: sort}, page, 3)
if pageErr != nil {
t.Fatal(pageErr)
}
if total != len(ids) {
t.Fatalf("page %d total=%d, want %d", page, total, len(ids))
}
for _, item := range items {
got = append(got, item.ID)
}
}
if !slices.Equal(got, want) {
t.Fatalf("same-timestamp pagination was unstable: got=%v want=%v", got, want)
}
})
}
}
func TestFindingGroupsAndUnassignedPaging(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
vc := fmt.Sprintf("__test_finding_groups_%d__", time.Now().UnixNano())
cleanupFindings := func() { _, _ = d.Exec(`DELETE FROM findings WHERE vulnclass=$1`, vc) }
defer cleanupFindings()
taskA, err := d.CreateTask("group A", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskA.ID)
taskB, err := d.CreateTask("group B", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(taskB.ID)
seed := []struct {
taskID int64
severity string
status string
}{
{taskA.ID, SeverityCritical, FindingPending},
{taskA.ID, SeverityHigh, FindingResolved},
{taskB.ID, SeverityLow, FindingPending},
{0, SeverityMedium, FindingPending},
}
for i, item := range seed {
id, addErr := d.AddFinding(item.taskID, 0, vc, "", item.severity, fmt.Sprintf("summary %d", i), "", "test", nil)
if addErr != nil {
t.Fatalf("AddFinding[%d]: %v", i, addErr)
}
if _, setErr := d.SetFindingStatus(id, item.status); setErr != nil {
t.Fatalf("SetFindingStatus[%d]: %v", i, setErr)
}
}
// Retained findings from deleted tasks join the same bucket as findings that
// were created without any task.
if err := d.DeleteTask(taskB.ID); err != nil {
t.Fatal(err)
}
page, groupTotal, findingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc, Sort: "severity"}, 1, 1)
if err != nil {
t.Fatal(err)
}
if len(page) != 1 || groupTotal != 2 || findingTotal != 4 {
t.Fatalf("page totals: items=%d groups=%d findings=%d", len(page), groupTotal, findingTotal)
}
groups, _, _, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
var live, unassigned *FindingGroup
for i := range groups {
if groups[i].TaskID == nil {
unassigned = &groups[i]
} else if *groups[i].TaskID == taskA.ID {
live = &groups[i]
}
}
if live == nil || live.Count != 2 || live.Critical != 1 || live.High != 1 || live.TaskDescription != "group A" {
t.Fatalf("live group mismatch: %+v", live)
}
if unassigned == nil || unassigned.Count != 2 || unassigned.Medium != 1 || unassigned.Low != 1 {
t.Fatalf("unassigned group mismatch: %+v", unassigned)
}
groupForTask := func(items []FindingGroup, taskID int64) *FindingGroup {
for i := range items {
if items[i].TaskID != nil && *items[i].TaskID == taskID {
return &items[i]
}
}
return nil
}
if err := d.SetPaused(taskA.ID, true); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "paused" {
t.Fatalf("paused task group status mismatch: %+v", groups)
}
if err := d.SetPaused(taskA.ID, false); err != nil {
t.Fatal(err)
}
if err := d.Enqueue(taskA.ID, "resume"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "queued" {
t.Fatalf("queued task group status mismatch: %+v", groups)
}
if err := d.SetStatus(taskA.ID, "done"); err != nil {
t.Fatal(err)
}
groups, _, _, err = d.ListFindingGroups(FindingFilter{VulnClass: vc}, 1, 10)
if err != nil {
t.Fatal(err)
}
if group := groupForTask(groups, taskA.ID); group == nil || group.TaskStatus != "done" {
t.Fatalf("terminal task status must win over queued flag: %+v", groups)
}
orphans, total, err := d.ListFindingsPage(FindingFilter{VulnClass: vc, TaskID: FindingUnassignedTask}, 1, 10)
if err != nil || total != 2 || len(orphans) != 2 {
t.Fatalf("unassigned page: len=%d total=%d err=%v", len(orphans), total, err)
}
maxPage := int(^uint(0) >> 1)
farFindings, farTotal, err := d.ListFindingsPage(FindingFilter{VulnClass: vc}, maxPage, 200)
if err != nil || farTotal != len(seed) || len(farFindings) != 0 {
t.Fatalf("far finding page: len=%d total=%d err=%v", len(farFindings), farTotal, err)
}
farGroups, farGroupTotal, farFindingTotal, err := d.ListFindingGroups(FindingFilter{VulnClass: vc}, maxPage, 100)
if err != nil || farGroupTotal != 2 || farFindingTotal != len(seed) || len(farGroups) != 0 {
t.Fatalf("far group page: len=%d groups=%d findings=%d err=%v",
len(farGroups), farGroupTotal, farFindingTotal, err)
}
filtered, filteredGroups, filteredFindings, err := d.ListFindingGroups(
FindingFilter{VulnClass: vc, Status: FindingResolved}, 1, 10,
)
if err != nil || filteredGroups != 1 || filteredFindings != 1 || len(filtered) != 1 || filtered[0].High != 1 {
t.Fatalf("filtered groups: %+v groups=%d findings=%d err=%v", filtered, filteredGroups, filteredFindings, err)
}
}
func TestAddFindingFollowUpIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
task, err := d.CreateTask("finding follow-up", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("follow-up-%d.example.test", time.Now().UnixNano()),
TaskID: task.ID,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id=$1`, assetID)
store := d.Exploration(task.ExplorationID)
findingNodeID, err := store.AddNode(KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
findingID, err := d.AddFinding(task.ID, findingNodeID, "test", "source", SeverityHigh, "source", "", "worker", []int64{assetID})
if err != nil {
t.Fatal(err)
}
auditInput := Activity{Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "验证可利用性并形成证据链"}
intentID, audit, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "验证可利用性并形成证据链", auditInput)
if err != nil {
t.Fatal(err)
}
secondID, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "从另一条路径深入", Activity{
Worker: "system", Kind: "text", Summary: "人工提交漏洞深入利用意图", Detail: "从另一条路径深入",
})
if err != nil {
t.Fatal(err)
}
if secondID == intentID {
t.Fatal("repeated follow-up submissions must create distinct intents")
}
expectedAuditSummary := fmt.Sprintf("%s #%d", auditInput.Summary, intentID)
if audit.ID <= 0 || audit.NodeID == nil || *audit.NodeID != intentID || audit.CreatedAt.IsZero() || audit.Summary != expectedAuditSummary {
t.Fatalf("persisted audit mismatch: %+v", audit)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil {
t.Fatalf("GetNode: node=%+v err=%v", node, err)
}
if node.Kind != KindIntent || node.Priority != 10 || node.State != "open" || node.Origin != "human" {
t.Fatalf("follow-up intent metadata: %+v", node)
}
var anchorCount, edgeCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_anchors WHERE node_id=$1 AND asset_id=$2`, intentID, assetID).Scan(&anchorCount); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_edges
WHERE exploration_id=$1 AND src_id=$2 AND rel=$3 AND dst_id=$4`,
task.ExplorationID, findingNodeID, RelDerivedFrom, intentID).Scan(&edgeCount); err != nil {
t.Fatal(err)
}
if anchorCount != 1 || edgeCount != 1 {
t.Fatalf("follow-up lineage: anchors=%d edges=%d", anchorCount, edgeCount)
}
var persistedAuditCount int
if err := d.QueryRow(`SELECT COUNT(*) FROM activity
WHERE id=$1 AND exploration_id=$2 AND node_id=$3 AND worker='system' AND summary=$4`,
audit.ID, task.ExplorationID, intentID, expectedAuditSummary).Scan(&persistedAuditCount); err != nil {
t.Fatal(err)
}
if persistedAuditCount != 1 {
t.Fatalf("atomic audit count=%d, want 1", persistedAuditCount)
}
var nodesBefore, activitiesBefore int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesBefore); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesBefore); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "must roll back", Activity{
Worker: "system", Kind: "text", Summary: "invalid audit", Metadata: json.RawMessage(`{`),
}); err == nil {
t.Fatal("invalid activity metadata should fail the transaction")
}
var nodesAfter, activitiesAfter int
if err := d.QueryRow(`SELECT COUNT(*) FROM exploration_nodes WHERE exploration_id=$1`, task.ExplorationID).Scan(&nodesAfter); err != nil {
t.Fatal(err)
}
if err := d.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.ExplorationID).Scan(&activitiesAfter); err != nil {
t.Fatal(err)
}
if nodesAfter != nodesBefore || activitiesAfter != activitiesBefore {
t.Fatalf("activity failure did not roll back graph: nodes %d->%d activities %d->%d",
nodesBefore, nodesAfter, activitiesBefore, activitiesAfter)
}
if _, err := d.Exec(`DELETE FROM exploration_nodes WHERE id=$1`, findingNodeID); err != nil {
t.Fatal(err)
}
if _, _, err := store.AddFindingFollowUpIntent(findingID, findingNodeID, "should fail", auditInput); !errors.Is(err, ErrFindingOriginUnavailable) {
t.Fatalf("missing source node: want ErrFindingOriginUnavailable, got %v", err)
}
}
+32
View File
@@ -0,0 +1,32 @@
package db
import "fmt"
// DiscardOpenIntent compensates a follow-up creation when task admission fails.
// The task execution gate must still be held by the caller, so the intent cannot
// be claimed between this check and deletion. Edges and anchors cascade with the
// node; activity is deleted explicitly because its node FK otherwise becomes NULL.
func (s *ExplorationStore) DiscardOpenIntent(id int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM activity WHERE exploration_id=$1 AND node_id=$2`, s.expID, id); err != nil {
return err
}
res, err := tx.Exec(`DELETE FROM exploration_nodes
WHERE id=$1 AND exploration_id=$2 AND kind='intent' AND state='open'`, id, s.expID)
if err != nil {
return err
}
removed, err := res.RowsAffected()
if err != nil {
return err
}
if removed != 1 {
return fmt.Errorf("open intent %d was not available for admission rollback", id)
}
return tx.Commit()
}
+285
View File
@@ -0,0 +1,285 @@
package db
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestCompareAndSetIntentStateAllowsSingleWinner(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("intent CAS", "only one controller wins")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "controlled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "worker"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
const contenders = 12
var winners atomic.Int32
var wg sync.WaitGroup
start := make(chan struct{})
for range contenders {
wg.Add(1)
go func() {
defer wg.Done()
<-start
changed, transitionErr := store.CompareAndSetIntentState(intentID, "running", "paused")
if transitionErr != nil {
t.Errorf("transition: %v", transitionErr)
return
}
if changed {
winners.Add(1)
}
}()
}
close(start)
wg.Wait()
if got := winners.Load(); got != 1 {
t.Fatalf("CAS winners=%d, want 1", got)
}
node, err := store.GetNode(intentID)
if err != nil || node == nil || node.State != "paused" {
t.Fatalf("node=%+v err=%v, want paused", node, err)
}
}
func TestCancelIntentPreservesTokenMeteringWithoutDoubleCount(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("cancel token rollup", "preserve consumed tokens")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
intentID, err := store.AddIntent(map[string]any{"summary": "cancelled"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if claimed, err := store.ClaimIntent(intentID, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
baselineDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
baseline := dailyTokenBuckets(baselineDaily)
appendUsage := func(kind string, at time.Time, input, output, read, write *int) {
t.Helper()
activityID, appendErr := store.AppendActivity(Activity{
NodeID: &intentID, Worker: "work#1", Kind: kind,
InputTokens: input, OutputTokens: output,
CacheReadTokens: read, CacheWriteTokens: write,
})
if appendErr != nil {
t.Fatal(appendErr)
}
if _, updateErr := d.Exec(`UPDATE activity SET created_at=$1 WHERE id=$2`, at, activityID); updateErr != nil {
t.Fatal(updateErr)
}
}
values := func(input, output, read, write int) (*int, *int, *int, *int) {
return &input, &output, &read, &write
}
now := time.Now().UTC()
atUTCNoon := func(daysAgo int) time.Time {
at := now.AddDate(0, 0, -daysAgo)
return time.Date(at.Year(), at.Month(), at.Day(), 12, 0, 0, 0, time.UTC)
}
authoritativeDay := atUTCNoon(6)
fallbackDay := atUTCNoon(4)
trailingDay := atUTCNoon(2)
// Run 1 has an authoritative result; its preceding cumulative frame must not
// be counted a second time or move usage to the frame's earlier date.
i, o, r, w := values(100, 100, 100, 100)
appendUsage("usage", atUTCNoon(7), i, o, r, w)
i, o, r, w = values(10, 11, 12, 13)
appendUsage("result", authoritativeDay, i, o, r, w)
// Run 2 models a legacy failed result that omitted usage. Its fallback belongs
// to the result's date, not the cumulative frame's date.
i, o, r, w = values(20, 21, 22, 23)
appendUsage("usage", atUTCNoon(5), i, o, r, w)
appendUsage("result", fallbackDay, nil, nil, nil, nil)
// Run 3 was interrupted before a terminal result was persisted.
i, o, r, w = values(30, 31, 32, 33)
appendUsage("usage", trailingDay, i, o, r, w)
beforeCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
beforeCancel := dailyTokenBuckets(beforeCancelDaily)
assertDailyTokenDelta(t, beforeCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, beforeCancel, baseline, fallbackDay, 0, 0, 0)
assertDailyTokenDelta(t, beforeCancel, baseline, trailingDay, 0, 0, 0)
if err := store.SetIntentState(intentID, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(intentID)
if err != nil {
t.Fatal(err)
}
if cleanup.Activities != 5 {
t.Fatalf("deleted activities=%d, want 5", cleanup.Activities)
}
total, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, total, 60, 63, 66, 69)
afterCancelDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
afterCancel := dailyTokenBuckets(afterCancelDaily)
assertDailyTokenDelta(t, afterCancel, baseline, authoritativeDay, 10, 11, 12)
assertDailyTokenDelta(t, afterCancel, baseline, fallbackDay, 20, 21, 22)
assertDailyTokenDelta(t, afterCancel, baseline, trailingDay, 30, 31, 32)
assertDailyTokenDelta(t, afterCancel, baseline, now, 0, 0, 0)
sessions, err := store.TokenStatsBySession()
if err != nil {
t.Fatal(err)
}
if len(sessions) != 0 {
t.Fatalf("cancelled intent leaked into sessions: %+v", sessions)
}
var rollups, datedRollups int
if err := d.QueryRow(`SELECT COUNT(*), COUNT(*) FILTER (
WHERE metadata->>'token_day'=TO_CHAR(created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD')
) FROM activity
WHERE exploration_id=$1 AND worker='token-ledger' AND kind='result'
AND metadata->>'cancelled_intent_id'=$2`, expID, fmt.Sprint(intentID)).Scan(&rollups, &datedRollups); err != nil {
t.Fatal(err)
}
if rollups != 3 || datedRollups != rollups {
t.Fatalf("token rollups=%d dated=%d, want three correctly dated rows", rollups, datedRollups)
}
if _, err := store.CancelIntent(intentID); err == nil {
t.Fatal("second cancellation unexpectedly succeeded")
}
afterRetry, err := store.TokenTotal()
if err != nil {
t.Fatal(err)
}
assertTokenUsage(t, afterRetry, 60, 63, 66, 69)
afterRetryDaily, err := d.TokenDailyAll(30)
if err != nil {
t.Fatal(err)
}
assertDailyTokenBucketsEqual(t, dailyTokenBuckets(afterRetryDaily), afterCancel)
}
func dailyTokenBuckets(items []DailyTokenBucket) map[string]DailyTokenBucket {
out := make(map[string]DailyTokenBucket, len(items))
for _, item := range items {
out[item.Day] = item
}
return out
}
func assertDailyTokenDelta(t *testing.T, got, baseline map[string]DailyTokenBucket, at time.Time, input, output, read int) {
t.Helper()
day := at.UTC().Format(time.DateOnly)
actual, initial := got[day], baseline[day]
if actual.InputTokens-initial.InputTokens != input ||
actual.OutputTokens-initial.OutputTokens != output ||
actual.CacheReadTokens-initial.CacheReadTokens != read {
t.Fatalf("token delta for %s = input:%d output:%d cache-read:%d, want %d/%d/%d",
day, actual.InputTokens-initial.InputTokens, actual.OutputTokens-initial.OutputTokens,
actual.CacheReadTokens-initial.CacheReadTokens, input, output, read)
}
}
func assertDailyTokenBucketsEqual(t *testing.T, got, want map[string]DailyTokenBucket) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("daily token buckets changed after retry: got=%+v want=%+v", got, want)
}
for day, expected := range want {
if actual, ok := got[day]; !ok || actual != expected {
t.Fatalf("daily token bucket %s changed after retry: got=%+v want=%+v", day, actual, expected)
}
}
}
func TestCancelIntentPreservesOutputsYieldedByAnotherIntent(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("shared intent output", "preserve shared facts and findings")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID)
store := d.Exploration(expID)
first, err := store.AddIntent(map[string]any{"summary": "first"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
second, err := store.AddIntent(map[string]any{"summary": "second"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
fact, err := store.AddNode(KindFact, map[string]any{"text": "shared fact"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
finding, err := store.AddNode(KindFinding, map[string]any{"summary": "shared finding"}, 1, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
for _, intentID := range []int64{first, second} {
if err := store.Link(intentID, RelYields, fact); err != nil {
t.Fatal(err)
}
if err := store.Link(intentID, RelYields, finding); err != nil {
t.Fatal(err)
}
}
if claimed, err := store.ClaimIntent(first, "work#1"); err != nil || !claimed {
t.Fatalf("claim: claimed=%v err=%v", claimed, err)
}
if err := store.SetIntentState(first, "paused"); err != nil {
t.Fatal(err)
}
cleanup, err := store.CancelIntent(first)
if err != nil {
t.Fatal(err)
}
if cleanup.Intents != 1 || cleanup.Facts != 0 || cleanup.Findings != 0 {
t.Fatalf("cleanup=%+v, want only the cancelled intent", cleanup)
}
for _, nodeID := range []int64{fact, finding} {
node, getErr := store.GetNode(nodeID)
if getErr != nil || node == nil {
t.Fatalf("shared node %d was removed: node=%+v err=%v", nodeID, node, getErr)
}
}
}
+59 -2
View File
@@ -31,6 +31,11 @@ CREATE TABLE IF NOT EXISTS assets (
'root_domain','ip','subdomain','app','service','endpoint'
)),
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
-- explicit: caller/user selected the company; scope: derived from company_scope.
-- Existing installations are conservatively migrated as explicit so a scope
-- rebuild can never erase a historical manual association.
company_source TEXT NOT NULL DEFAULT 'explicit'
CHECK (company_source IN ('explicit','scope')),
task_ids BIGINT[] NOT NULL DEFAULT '{}',
domain TEXT,
root_domain TEXT,
@@ -84,6 +89,13 @@ CREATE INDEX IF NOT EXISTS idx_av2_bound_domains ON assets USING GIN(bound_domai
CREATE INDEX IF NOT EXISTS idx_av2_open_ports ON assets USING GIN(open_ports) WHERE type = 'ip';
CREATE INDEX IF NOT EXISTS idx_av2_last_seen ON assets(last_seen DESC);
CREATE INDEX IF NOT EXISTS idx_av2_type_seen ON assets(type, last_seen DESC);
ALTER TABLE assets ADD COLUMN IF NOT EXISTS company_source TEXT;
UPDATE assets SET company_source = 'explicit' WHERE company_source IS NULL;
ALTER TABLE assets ALTER COLUMN company_source SET DEFAULT 'explicit';
ALTER TABLE assets ALTER COLUMN company_source SET NOT NULL;
ALTER TABLE assets DROP CONSTRAINT IF EXISTS assets_company_source_check;
ALTER TABLE assets ADD CONSTRAINT assets_company_source_check
CHECK (company_source IN ('explicit','scope'));
DROP TRIGGER IF EXISTS trg_av2_upd ON assets;
CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
@@ -91,17 +103,36 @@ CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
CREATE TABLE IF NOT EXISTS company_scope (
id BIGSERIAL PRIMARY KEY,
company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr')),
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr','icp','keyword')),
domain TEXT,
net CIDR,
value TEXT,
raw TEXT NOT NULL,
reason TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT uq_sv2_domain UNIQUE (company_id, domain),
CONSTRAINT uq_sv2_net UNIQUE (company_id, net)
CONSTRAINT uq_sv2_net UNIQUE (company_id, net),
CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
)
);
-- Existing installations need the new text payload and expanded kind check.
ALTER TABLE company_scope ADD COLUMN IF NOT EXISTS value TEXT;
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS company_scope_kind_check;
ALTER TABLE company_scope ADD CONSTRAINT company_scope_kind_check
CHECK (kind IN ('domain','ip','cidr','icp','keyword'));
ALTER TABLE company_scope DROP CONSTRAINT IF EXISTS ck_company_scope_payload;
ALTER TABLE company_scope ADD CONSTRAINT ck_company_scope_payload CHECK (
(kind = 'domain' AND domain IS NOT NULL AND net IS NULL AND value IS NULL)
OR (kind IN ('ip','cidr') AND domain IS NULL AND net IS NOT NULL AND value IS NULL)
OR (kind IN ('icp','keyword') AND domain IS NULL AND net IS NULL AND value IS NOT NULL)
);
CREATE INDEX IF NOT EXISTS idx_sv2_domain ON company_scope(domain) WHERE kind = 'domain';
CREATE INDEX IF NOT EXISTS idx_sv2_net ON company_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr');
CREATE UNIQUE INDEX IF NOT EXISTS uq_sv2_value ON company_scope(company_id, kind, value) WHERE kind IN ('icp','keyword');
CREATE INDEX IF NOT EXISTS idx_sv2_icp ON company_scope(value) WHERE kind = 'icp';
CREATE INDEX IF NOT EXISTS idx_sv2_company ON company_scope(company_id);
-- =====================================================================
@@ -301,6 +332,8 @@ CREATE TABLE IF NOT EXISTS tasks (
CHECK (status IN ('created','running','paused','done','failed','timeout')),
paused BOOLEAN NOT NULL DEFAULT false,
queued BOOLEAN NOT NULL DEFAULT false,
queued_at TIMESTAMPTZ,
queue_mode TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
llm_chain_revision BIGINT NOT NULL DEFAULT 0,
@@ -324,9 +357,30 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300;
-- 并发上限挂起态;补旧库。true=因并发上限排队、等待空位自动启动。
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false;
-- queued_at makes admission FIFO reflect the actual enqueue order rather than the
-- task creation order. queue_mode distinguishes first bootstrap from resuming an
-- exploration that already owns goals/history.
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued_at TIMESTAMPTZ;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queue_mode TEXT NOT NULL DEFAULT '';
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS active_llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_chain_revision BIGINT NOT NULL DEFAULT 0;
-- Reusable task description/goal presets. nkey is the normalized, case-insensitive
-- identity used to reject visually equivalent duplicate names.
CREATE TABLE IF NOT EXISTS task_templates (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
nkey TEXT NOT NULL UNIQUE,
description TEXT NOT NULL,
goal TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_task_templates_updated ON task_templates(updated_at DESC, id DESC);
DROP TRIGGER IF EXISTS trg_task_templates_upd ON task_templates;
CREATE TRIGGER trg_task_templates_upd BEFORE UPDATE ON task_templates
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- Direct, read-only task context inheritance. Relations are intentionally not
-- recursive: a task sees only the source tasks explicitly chosen at creation.
CREATE TABLE IF NOT EXISTS task_relations (
@@ -573,10 +627,13 @@ CREATE TABLE IF NOT EXISTS conversations (
agent_key TEXT NOT NULL,
title TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
pinned_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE conversations ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ;
CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned_at DESC) WHERE pinned_at IS NOT NULL;
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
+33 -2
View File
@@ -60,6 +60,11 @@ func (d *DB) hydrateTaskContext(t *Task) error {
return err
}
t.SourceTaskIDs = sources
companies, err := d.TaskCompanyIDs(t.ID)
if err != nil {
return err
}
t.CompanyIDs = companies
t.LLMProfileIDs = make([]int64, 0, len(chain))
t.LLMFailoverState = "default"
var latest *TaskLLMProfile
@@ -172,6 +177,28 @@ func (d *DB) TaskSourceIDs(taskID int64) ([]int64, error) {
return out, rows.Err()
}
// TaskCompanyIDs returns the companies whose asset scopes are available to the
// task. The task_scope rows remain the single source of truth.
func (d *DB) TaskCompanyIDs(taskID int64) ([]int64, error) {
rows, err := d.Query(`
SELECT company_id FROM task_scope
WHERE task_id=$1 AND kind='company' AND company_id IS NOT NULL
ORDER BY id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
func (d *DB) TaskSources(taskID int64) ([]TaskSource, error) {
rows, err := d.Query(`
SELECT t.id, t.exploration_id, t.description, t.goal, t.status
@@ -339,13 +366,17 @@ ORDER BY position LIMIT 1`, taskID, position).Scan(&next)
case nil:
out.Advanced = true
out.NextProfileID = &next
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=$2, llm_profile_id=$2 WHERE id=$1`, taskID, next); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID, next); err != nil {
return out, err
}
case sql.ErrNoRows:
out.Advanced = true
out.ChainExhausted = true
if _, err := tx.Exec(`UPDATE tasks SET active_llm_profile_id=NULL, llm_profile_id=NULL WHERE id=$1`, taskID); err != nil {
if _, err := tx.Exec(`UPDATE tasks
SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, taskID); err != nil {
return out, err
}
default:
+66
View File
@@ -0,0 +1,66 @@
package db
import (
"testing"
"time"
)
func TestTaskQueuePreservesBootstrapAndFIFOPosition(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("queue metadata", "keep first-run mode", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer d.DeleteTask(task.ID)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
first, err := d.GetTask(task.ID)
if err != nil || first == nil || first.QueuedAt == nil {
t.Fatalf("first enqueue: task=%+v err=%v", first, err)
}
// A follow-up/rerun can try to admit an already queued task as resume. The
// original bootstrap mode and FIFO timestamp must remain authoritative.
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "resume"); err != nil {
t.Fatal(err)
}
second, err := d.GetTask(task.ID)
if err != nil || second == nil {
t.Fatalf("second enqueue: task=%+v err=%v", second, err)
}
if second.QueueMode != "bootstrap" {
t.Fatalf("queue mode=%q, want bootstrap", second.QueueMode)
}
if second.QueuedAt == nil || !second.QueuedAt.Equal(*first.QueuedAt) {
t.Fatalf("repeated enqueue moved FIFO position: first=%v second=%v", first.QueuedAt, second.QueuedAt)
}
// Pausing a queued task removes it from the queue but keeps its required
// startup mode. A later requeue receives a new tail position.
if err := d.Dequeue(task.ID, false); err != nil {
t.Fatal(err)
}
paused, err := d.GetTask(task.ID)
if err != nil || paused == nil || paused.Queued || paused.QueueMode != "bootstrap" {
t.Fatalf("paused queue metadata: task=%+v err=%v", paused, err)
}
time.Sleep(time.Millisecond)
if err := d.Enqueue(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
requeued, err := d.GetTask(task.ID)
if err != nil || requeued == nil || requeued.QueuedAt == nil {
t.Fatalf("requeue: task=%+v err=%v", requeued, err)
}
if !requeued.QueuedAt.After(*first.QueuedAt) {
t.Fatalf("requeue did not move to FIFO tail: first=%v requeued=%v", first.QueuedAt, requeued.QueuedAt)
}
}
+27
View File
@@ -1,6 +1,8 @@
package db
import (
"errors"
"reflect"
"strings"
"testing"
)
@@ -18,3 +20,28 @@ func TestCreateTaskRejectsTooManySourcesBeforeOpeningTransaction(t *testing.T) {
t.Fatalf("expected source-count validation error, got %v", err)
}
}
func TestNormalizeTaskCompanyIDs(t *testing.T) {
got, err := NormalizeTaskCompanyIDs([]int64{4, 2, 4, 7, 2})
if err != nil {
t.Fatal(err)
}
if want := []int64{4, 2, 7}; !reflect.DeepEqual(got, want) {
t.Fatalf("NormalizeTaskCompanyIDs=%v, want %v", got, want)
}
if _, err := NormalizeTaskCompanyIDs([]int64{1, 0}); !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("invalid company id error=%v", err)
}
}
func TestCreateTaskRejectsTooManyCompaniesBeforeOpeningTransaction(t *testing.T) {
companyIDs := make([]int64, MaxTaskCompanyCount+1)
for i := range companyIDs {
companyIDs[i] = int64(i + 1)
}
_, err := (&DB{}).CreateTaskWithOptions("child", "goal", TaskCreateOptions{CompanyIDs: companyIDs})
if !errors.Is(err, ErrTaskCompanyIDsInvalid) {
t.Fatalf("expected company-count validation error, got %v", err)
}
}
+231
View File
@@ -0,0 +1,231 @@
package db
import (
"database/sql"
"errors"
"fmt"
"strings"
"time"
"unicode/utf8"
"github.com/jackc/pgx/v5/pgconn"
)
const (
MaxTaskTemplateNameRunes = 120
MaxTaskTemplateTextRunes = 16000
)
var (
ErrTaskTemplateInvalid = errors.New("invalid task template")
ErrTaskTemplateNameConflict = errors.New("task template name already exists")
ErrTaskTemplateNotFound = errors.New("task template not found")
)
// TaskTemplate is a reusable task description and goal preset.
type TaskTemplate struct {
ID int64 `json:"id"`
Name string `json:"name"`
NKey string `json:"-"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// TaskTemplateInput is the create/update payload after normalization.
type TaskTemplateInput struct {
Name string
Description string
Goal string
}
// TaskTemplatePatch changes only fields whose pointers are non-nil.
type TaskTemplatePatch struct {
Name *string
Description *string
Goal *string
}
const taskTemplateCols = `id, name, nkey, description, goal, created_at, updated_at`
func scanTaskTemplate(row interface{ Scan(...any) error }) (TaskTemplate, error) {
var t TaskTemplate
err := row.Scan(&t.ID, &t.Name, &t.NKey, &t.Description, &t.Goal, &t.CreatedAt, &t.UpdatedAt)
return t, err
}
// taskTemplateName normalizes display whitespace while preserving the user's case.
func taskTemplateName(name string) string { return strings.Join(strings.Fields(name), " ") }
// taskTemplateNKey is the case-insensitive identity used by the unique index.
func taskTemplateNKey(name string) string { return strings.ToLower(taskTemplateName(name)) }
func normalizeTaskTemplateInput(in TaskTemplateInput) (TaskTemplateInput, string, error) {
in.Name = taskTemplateName(in.Name)
in.Description = strings.TrimSpace(in.Description)
in.Goal = strings.TrimSpace(in.Goal)
switch {
case in.Name == "":
return in, "", fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Name) > MaxTaskTemplateNameRunes:
return in, "", fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
case in.Description == "":
return in, "", fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Description) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
case in.Goal == "":
return in, "", fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
case utf8.RuneCountInString(in.Goal) > MaxTaskTemplateTextRunes:
return in, "", fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
return in, taskTemplateNKey(in.Name), nil
}
func normalizeTaskTemplatePatch(patch TaskTemplatePatch) (TaskTemplatePatch, *string, error) {
if patch.Name == nil && patch.Description == nil && patch.Goal == nil {
return patch, nil, fmt.Errorf("%w: no fields supplied", ErrTaskTemplateInvalid)
}
var nkey *string
if patch.Name != nil {
name := taskTemplateName(*patch.Name)
if name == "" {
return patch, nil, fmt.Errorf("%w: name is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(name) > MaxTaskTemplateNameRunes {
return patch, nil, fmt.Errorf("%w: name exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateNameRunes)
}
key := taskTemplateNKey(name)
patch.Name = &name
nkey = &key
}
if patch.Description != nil {
description := strings.TrimSpace(*patch.Description)
if description == "" {
return patch, nil, fmt.Errorf("%w: description is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(description) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: description exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Description = &description
}
if patch.Goal != nil {
goal := strings.TrimSpace(*patch.Goal)
if goal == "" {
return patch, nil, fmt.Errorf("%w: goal is required", ErrTaskTemplateInvalid)
}
if utf8.RuneCountInString(goal) > MaxTaskTemplateTextRunes {
return patch, nil, fmt.Errorf("%w: goal exceeds %d characters", ErrTaskTemplateInvalid, MaxTaskTemplateTextRunes)
}
patch.Goal = &goal
}
return patch, nkey, nil
}
func taskTemplateUniqueViolation(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "23505"
}
// CreateTaskTemplate inserts one globally reusable preset.
func (d *DB) CreateTaskTemplate(in TaskTemplateInput) (*TaskTemplate, error) {
in, nkey, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`
INSERT INTO task_templates(name, nkey, description, goal)
VALUES ($1,$2,$3,$4)
ON CONFLICT (nkey) DO NOTHING
RETURNING `+taskTemplateCols, in.Name, nkey, in.Description, in.Goal))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// ListTaskTemplates returns the most recently maintained templates first.
func (d *DB) ListTaskTemplates() ([]*TaskTemplate, error) {
rows, err := d.Query(`SELECT ` + taskTemplateCols + ` FROM task_templates ORDER BY updated_at DESC, id DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*TaskTemplate{}
for rows.Next() {
t, err := scanTaskTemplate(rows)
if err != nil {
return nil, err
}
out = append(out, &t)
}
return out, rows.Err()
}
// GetTaskTemplate returns nil when id does not exist.
func (d *DB) GetTaskTemplate(id int64) (*TaskTemplate, error) {
t, err := scanTaskTemplate(d.QueryRow(`SELECT `+taskTemplateCols+` FROM task_templates WHERE id=$1`, id))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &t, nil
}
// UpdateTaskTemplate replaces the editable fields of one preset.
func (d *DB) UpdateTaskTemplate(id int64, in TaskTemplateInput) (*TaskTemplate, error) {
in, _, err := normalizeTaskTemplateInput(in)
if err != nil {
return nil, err
}
return d.PatchTaskTemplate(id, TaskTemplatePatch{
Name: &in.Name, Description: &in.Description, Goal: &in.Goal,
})
}
// PatchTaskTemplate atomically changes only the supplied fields. Keeping the
// merge in one UPDATE prevents concurrent disjoint PATCH requests from losing
// each other's changes.
func (d *DB) PatchTaskTemplate(id int64, patch TaskTemplatePatch) (*TaskTemplate, error) {
patch, nkey, err := normalizeTaskTemplatePatch(patch)
if err != nil {
return nil, err
}
t, err := scanTaskTemplate(d.QueryRow(`UPDATE task_templates
SET name=CASE WHEN $2 THEN $3::text ELSE name END,
nkey=CASE WHEN $2 THEN $4::text ELSE nkey END,
description=CASE WHEN $5 THEN $6::text ELSE description END,
goal=CASE WHEN $7 THEN $8::text ELSE goal END
WHERE id=$1
RETURNING `+taskTemplateCols,
id,
patch.Name != nil, patch.Name, nkey,
patch.Description != nil, patch.Description,
patch.Goal != nil, patch.Goal,
))
if err == sql.ErrNoRows {
return nil, ErrTaskTemplateNotFound
}
if taskTemplateUniqueViolation(err) {
return nil, ErrTaskTemplateNameConflict
}
if err != nil {
return nil, err
}
return &t, nil
}
// DeleteTaskTemplate deletes one preset and reports whether it existed.
func (d *DB) DeleteTaskTemplate(id int64) (bool, error) {
result, err := d.Exec(`DELETE FROM task_templates WHERE id=$1`, id)
if err != nil {
return false, err
}
n, err := result.RowsAffected()
return n > 0, err
}
+131
View File
@@ -0,0 +1,131 @@
package db
import (
"errors"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestTaskTemplateCRUDAndNormalizedUniqueness(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
name := fmt.Sprintf("Template %d", suffix)
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: " " + strings.ReplaceAll(name, " ", " ") + " ",
Description: " initial description ",
Goal: " initial goal ",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
if created.Name != name || created.Description != "initial description" || created.Goal != "initial goal" {
t.Fatalf("template was not normalized: %+v", created)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: strings.ToUpper(name), Description: "duplicate", Goal: "duplicate",
}); !errors.Is(err, ErrTaskTemplateNameConflict) {
t.Fatalf("duplicate create error = %v, want %v", err, ErrTaskTemplateNameConflict)
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil || got == nil || got.Name != name {
t.Fatalf("GetTaskTemplate = %+v, %v", got, err)
}
listed, err := d.ListTaskTemplates()
if err != nil {
t.Fatal(err)
}
found := false
for _, template := range listed {
if template.ID == created.ID {
found = true
break
}
}
if !found {
t.Fatalf("created template %d missing from list", created.ID)
}
updatedName := name + " updated"
updated, err := d.UpdateTaskTemplate(created.ID, TaskTemplateInput{
Name: updatedName, Description: "new description", Goal: "new goal",
})
if err != nil {
t.Fatal(err)
}
if updated.Name != updatedName || updated.Description != "new description" || updated.Goal != "new goal" {
t.Fatalf("unexpected updated template: %+v", updated)
}
if _, err := d.CreateTaskTemplate(TaskTemplateInput{Name: "", Description: "x", Goal: "y"}); !errors.Is(err, ErrTaskTemplateInvalid) {
t.Fatalf("empty name error = %v, want %v", err, ErrTaskTemplateInvalid)
}
deleted, err := d.DeleteTaskTemplate(created.ID)
if err != nil || !deleted {
t.Fatalf("DeleteTaskTemplate = %v, %v", deleted, err)
}
deleted, err = d.DeleteTaskTemplate(created.ID)
if err != nil || deleted {
t.Fatalf("second DeleteTaskTemplate = %v, %v", deleted, err)
}
}
func TestTaskTemplateDisjointPatchesCompose(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
created, err := d.CreateTaskTemplate(TaskTemplateInput{
Name: fmt.Sprintf("Concurrent template %d", time.Now().UnixNano()),
Description: "initial description",
Goal: "initial goal",
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
description := "description from concurrent patch"
goal := "goal from concurrent patch"
start := make(chan struct{})
errs := make(chan error, 2)
var wg sync.WaitGroup
for _, patch := range []TaskTemplatePatch{{Description: &description}, {Goal: &goal}} {
patch := patch
wg.Add(1)
go func() {
defer wg.Done()
<-start
_, err := d.PatchTaskTemplate(created.ID, patch)
errs <- err
}()
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatal(err)
}
}
got, err := d.GetTaskTemplate(created.ID)
if err != nil {
t.Fatal(err)
}
if got == nil || got.Description != description || got.Goal != goal {
t.Fatalf("disjoint patches lost an update: %+v", got)
}
}
+119 -12
View File
@@ -3,6 +3,7 @@ package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"strconv"
"time"
@@ -10,14 +11,16 @@ import (
// Task is a row in the task registry (1:1 with an exploration).
type Task struct {
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
QueuedAt *time.Time `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
// for ActiveLLMProfileID while older API clients still send one profile id.
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
@@ -26,6 +29,7 @@ type Task struct {
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
@@ -74,6 +78,40 @@ const MinPlanHeartbeatSeconds = 600
// multiplying graph and asset-context queries without limit.
const MaxTaskSourceCount = 8
// MaxTaskCompanyCount bounds the number of company asset scopes attached to a
// task. Company scopes are prompt context, so an unbounded list would inflate
// every planner and worker call.
const MaxTaskCompanyCount = 32
var (
ErrTaskCompanyIDsInvalid = errors.New("invalid task company ids")
ErrTaskCompanyNotFound = errors.New("task company not found")
)
// NormalizeTaskCompanyIDs validates IDs and removes duplicates while retaining
// the user's first-seen order.
func NormalizeTaskCompanyIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, nil
}
seen := make(map[int64]struct{}, min(len(ids), MaxTaskCompanyCount))
normalized := make([]int64, 0, min(len(ids), MaxTaskCompanyCount))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("%w: company id must be positive", ErrTaskCompanyIDsInvalid)
}
if _, exists := seen[id]; exists {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
if len(normalized) > MaxTaskCompanyCount {
return nil, fmt.Errorf("%w: got more than %d unique companies", ErrTaskCompanyIDsInvalid, MaxTaskCompanyCount)
}
}
return normalized, nil
}
func normalizeHeartbeat(sec int) int {
if sec < MinPlanHeartbeatSeconds {
return MinPlanHeartbeatSeconds
@@ -95,6 +133,7 @@ func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSe
// with the task/exploration row.
type TaskCreateOptions struct {
SourceTaskIDs []int64
CompanyIDs []int64
LLMProfileIDs []int64
TimeoutSeconds int
PlanHeartbeatSeconds int
@@ -106,6 +145,11 @@ func (d *DB) CreateTaskWithOptions(description, goal string, opts TaskCreateOpti
if len(opts.SourceTaskIDs) > MaxTaskSourceCount {
return nil, fmt.Errorf("too many source tasks: got %d, maximum is %d", len(opts.SourceTaskIDs), MaxTaskSourceCount)
}
companyIDs, err := NormalizeTaskCompanyIDs(opts.CompanyIDs)
if err != nil {
return nil, err
}
opts.CompanyIDs = companyIDs
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -144,6 +188,7 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload));
LLMProfileID: active, ActiveLLMProfileID: active,
LLMProfileIDs: append([]int64(nil), opts.LLMProfileIDs...),
SourceTaskIDs: append([]int64(nil), opts.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), opts.CompanyIDs...),
TimeoutSeconds: opts.TimeoutSeconds, PlanHeartbeatSeconds: opts.PlanHeartbeatSeconds,
}
if err := tx.QueryRow(`
@@ -155,6 +200,9 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil {
return nil, err
}
if err := insertTaskCompanies(tx, t.ID, opts.CompanyIDs); err != nil {
return nil, err
}
if err := insertTaskLLMProfiles(tx, t.ID, opts.LLMProfileIDs); err != nil {
return nil, err
}
@@ -166,6 +214,33 @@ RETURNING id, status, paused, created_at`, description, goal, expID, active, opt
return t, tx.Commit()
}
func insertTaskCompanies(tx *sql.Tx, taskID int64, companyIDs []int64) error {
if len(companyIDs) == 0 {
return nil
}
var inserted int
err := tx.QueryRow(`
WITH requested(company_id, position) AS (
SELECT company_id, position
FROM unnest($2::bigint[]) WITH ORDINALITY AS requested(company_id, position)
), inserted AS (
INSERT INTO task_scope(task_id, kind, company_id, source, reason)
SELECT $1, 'company', companies.id, 'manual', '任务创建时关联企业'
FROM requested
JOIN companies ON companies.id=requested.company_id
ORDER BY requested.position
RETURNING company_id
)
SELECT count(*) FROM inserted`, taskID, companyIDs).Scan(&inserted)
if err != nil {
return err
}
if inserted != len(companyIDs) {
return fmt.Errorf("%w: one or more companies do not exist", ErrTaskCompanyNotFound)
}
return nil
}
func insertTaskRelations(tx *sql.Tx, taskID int64, sourceIDs []int64) error {
seen := map[int64]bool{}
for _, sourceID := range sourceIDs {
@@ -199,11 +274,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
return nil
}
const taskCols = `id, description, goal, exploration_id, status, paused, queued, llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
const taskCols = `id, description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
return nil, err
}
return &t, nil
@@ -263,11 +338,43 @@ func (d *DB) SetPaused(id int64, paused bool) error {
return err
}
func (d *DB) SetQueued(id int64, queued bool) error {
_, err := d.Exec(`UPDATE tasks SET queued=$1 WHERE id=$2`, queued, id)
// Enqueue places a task at the tail of the persistent admission queue. Repeating
// the operation while it is already queued keeps its original FIFO position.
func (d *DB) Enqueue(id int64, mode string) error {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
_, err := d.Exec(`UPDATE tasks
SET queued=true,
queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END,
queue_mode=CASE
WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END
WHERE id=$1`, id, mode)
return err
}
// Dequeue removes the concurrency hold. clearMode is false when a user pauses a
// queued task, preserving whether its next admission must bootstrap or resume.
func (d *DB) Dequeue(id int64, clearMode bool) error {
_, err := d.Exec(`UPDATE tasks
SET queued=false,
queued_at=NULL,
queue_mode=CASE WHEN $2 THEN '' ELSE queue_mode END
WHERE id=$1`, id, clearMode)
return err
}
// SetQueued is the compatibility helper used by older callers and tests. New
// scheduling code should use Enqueue/Dequeue so FIFO metadata is explicit.
func (d *DB) SetQueued(id int64, queued bool) error {
if queued {
return d.Enqueue(id, "bootstrap")
}
return d.Dequeue(id, true)
}
// SetStatus updates a task's lifecycle status. Entering a terminal state
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
+80
View File
@@ -1,6 +1,7 @@
package db
import (
"errors"
"fmt"
"sync"
"testing"
@@ -277,6 +278,7 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if got.LLMChainRevision <= initialRevision || got.ActiveLLMProfileID == nil || *got.ActiveLLMProfileID != profileIDs[0] {
t.Fatalf("replacement revision/cursor not preserved: %+v", got)
}
replacementRevision := got.LLMChainRevision
transition, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "insufficient_quota")
if err != nil {
@@ -285,6 +287,13 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
if !transition.Advanced || transition.ChainExhausted || transition.NextProfileID == nil || *transition.NextProfileID != profileIDs[1] {
t.Fatalf("unexpected first transition: %+v", transition)
}
got, err = d.GetTask(child.ID)
if err != nil {
t.Fatal(err)
}
if got.LLMChainRevision != replacementRevision+1 {
t.Fatalf("automatic cursor advance did not increment revision: before=%d after=%d", replacementRevision, got.LLMChainRevision)
}
late, err := d.MarkTaskLLMProfileQuotaExhausted(child.ID, profileIDs[0], "late duplicate")
if err != nil {
t.Fatal(err)
@@ -483,3 +492,74 @@ func TestTaskContextRejectsDuplicatesAndTerminalEdits(t *testing.T) {
t.Fatal("terminal task LLM edit should be rejected")
}
}
func TestCreateTaskWithCompanyScopes(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
suffix := time.Now().UnixNano()
companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "")
if err != nil {
t.Fatal(err)
}
companyB, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company B %d", suffix), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, companyA, companyB)
})
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyB},
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("creation result company IDs=%v", task.CompanyIDs)
}
got, err := d.GetTask(task.ID)
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(got.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("hydrated company IDs=%v", got.CompanyIDs)
}
var scopeCount int
if err := d.QueryRow(`SELECT count(*) FROM task_scope WHERE task_id=$1 AND kind='company'`, task.ID).Scan(&scopeCount); err != nil {
t.Fatal(err)
}
if scopeCount != 2 {
t.Fatalf("company task scope rows=%d want 2", scopeCount)
}
duplicateTask, err := d.CreateTaskWithOptions("duplicate company scope", "deduplicate", TaskCreateOptions{
CompanyIDs: []int64{companyA, companyA, companyB, companyA},
})
if err != nil {
t.Fatal(err)
}
if fmt.Sprint(duplicateTask.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
t.Fatalf("company IDs were not normalized: %v", duplicateTask.CompanyIDs)
}
if err := d.DeleteTask(duplicateTask.ID); err != nil {
t.Fatal(err)
}
badDescription := fmt.Sprintf("invalid-company-%d", suffix)
if _, err := d.CreateTaskWithOptions(badDescription, "rollback", TaskCreateOptions{CompanyIDs: []int64{companyA, 1 << 62}}); !errors.Is(err, ErrTaskCompanyNotFound) {
t.Fatalf("missing company error=%v, want %v", err, ErrTaskCompanyNotFound)
}
var leaked int
if err := d.QueryRow(`SELECT count(*) FROM explorations WHERE description=$1`, badDescription).Scan(&leaked); err != nil {
t.Fatal(err)
}
if leaked != 0 {
t.Fatalf("failed company association leaked %d exploration rows", leaked)
}
}
+151 -52
View File
@@ -2,12 +2,58 @@ package server
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/db"
)
// companyScopeInputs accepts both the new [{kind,value}] contract and the
// historical ["example.com","203.0.113.10"] contract.
type companyScopeInputs []db.ScopeInput
const maxCompanyMutationBodyBytes = 2 << 20
func decodeCompanyMutationRequest(w http.ResponseWriter, r *http.Request, value any) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxCompanyMutationBodyBytes)
if err := json.NewDecoder(r.Body).Decode(value); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return false
}
return true
}
func (items *companyScopeInputs) UnmarshalJSON(data []byte) error {
var rawItems []json.RawMessage
if err := json.Unmarshal(data, &rawItems); err != nil {
return err
}
out := make([]db.ScopeInput, 0, len(rawItems))
for i, raw := range rawItems {
var legacy string
if err := json.Unmarshal(raw, &legacy); err == nil {
out = append(out, db.ScopeInput{Value: legacy})
continue
}
var structured db.ScopeInput
if err := json.Unmarshal(raw, &structured); err != nil {
return fmt.Errorf("scope[%d] must be a string or {kind,value}", i)
}
out = append(out, structured)
}
*items = out
return nil
}
// assetStore returns the asset store.
func (s *Server) assetStore() *db.AssetStore {
if s.m.pg == nil {
@@ -53,32 +99,45 @@ func (s *Server) createCompany(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Name string `json:"name"`
Logo string `json:"logo"`
Scope []string `json:"scope"`
Name string `json:"name"`
Logo string `json:"logo"`
Scope companyScopeInputs `json:"scope"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
req.Name = strings.TrimSpace(req.Name)
if req.Name == "" {
writeErr(w, 400, "name required")
return
}
id, created, err := cs.UpsertCompany(req.Name, req.Logo)
if err != nil {
writeErr(w, 500, err.Error())
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
out := map[string]any{"id": id, "created": created}
if len(req.Scope) > 0 {
added, skipped, invalid, errs := cs.AddScope(id, req.Scope, "api")
out["scope_added"] = added
out["scope_skipped"] = skipped
out["scope_invalid"] = invalid
if len(errs) > 0 {
out["scope_errors"] = errs
id, added, skipped, invalid, scopeErrs, err := cs.CreateCompanyWithScope(req.Name, req.Logo, req.Scope, "api")
if err != nil {
if errors.Is(err, db.ErrCompanyNameConflict) {
writeErr(w, http.StatusConflict, "企业名称已存在")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(err, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
out := map[string]any{
"id": id,
"created": true,
"scope_added": added,
"scope_skipped": skipped,
"scope_invalid": invalid,
}
if len(scopeErrs) > 0 {
out["scope_errors"] = scopeErrs
}
writeJSON(w, 201, out)
}
@@ -131,20 +190,37 @@ func (s *Server) addCompanyScope(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Scope []string `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
Scope companyScopeInputs `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
var added, skipped, invalid int
var errs []string
var mutationErr error
if req.Reset {
added, invalid, errs = cs.UpdateScope(id, req.Scope, req.Reason)
added, invalid, errs, mutationErr = cs.UpdateScopeInputsChecked(id, req.Scope, req.Reason)
} else {
added, skipped, invalid, errs = cs.AddScope(id, req.Scope, req.Reason)
added, skipped, invalid, errs, mutationErr = cs.AddScopeInputsChecked(id, req.Scope, req.Reason)
}
if mutationErr != nil {
if errors.Is(mutationErr, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(mutationErr, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, mutationErr.Error())
return
}
out := map[string]any{
"added": added,
@@ -175,21 +251,32 @@ func (s *Server) deleteCompany(w http.ResponseWriter, r *http.Request) {
var req struct {
DeleteAssets bool `json:"delete_assets"`
}
// Body is optional; ignore decode errors (e.g. empty body).
_ = json.NewDecoder(r.Body).Decode(&req)
assetsDeleted := int64(0)
if req.DeleteAssets {
as := s.assetStore()
if as != nil {
assetsDeleted, err = as.DeleteByCompanyID(id)
if err != nil {
writeErr(w, 500, err.Error())
return
// The body is optional. Only an actual empty body is ignored; malformed or
// trailing JSON is a client error.
decoder := json.NewDecoder(r.Body)
if err := decoder.Decode(&req); err != nil {
if !errors.Is(err, io.EOF) {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
return
}
} else {
var trailing any
if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
if err == nil {
writeErr(w, http.StatusBadRequest, "invalid JSON: multiple values")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return
}
}
if err := cs.DeleteCompany(id); err != nil {
assetsDeleted, err := s.m.DeleteCompanyWithAssets(id, req.DeleteAssets)
if err != nil {
if errors.Is(err, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
writeErr(w, 500, err.Error())
return
}
@@ -217,7 +304,10 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
// GET /api/assets
// =====================================================================
const defaultAssetPageSize = 50
const (
defaultAssetPageSize = 50
maxAssetPageSize = 200
)
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
as := s.assetStore()
@@ -231,46 +321,55 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
offset, _ := strconv.Atoi(q.Get("offset"))
if limit <= 0 {
limit = defaultAssetPageSize
} else if limit > maxAssetPageSize {
limit = maxAssetPageSize
}
if offset < 0 {
offset = 0
}
var assets []*db.Asset
assets := []*db.Asset{}
var err error
// total is the full match count ignoring limit/offset — lets the UI paginate on
// the server. For company/task scopes it falls back to the returned length.
// total is the full match count ignoring limit/offset. Count first so an offset
// beyond the last row can return an empty page without an expensive scan.
total := 0
if dsl := q.Get("dsl"); dsl != "" {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
if err == nil {
total, _ = as.CountDSL(dsl, typ)
if err := db.ValidateDSL(dsl); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
total, err = as.CountDSL(dsl, typ)
if err == nil && offset < total {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
}
} else {
companyID, _ := strconv.ParseInt(q.Get("company_id"), 10, 64)
taskID, _ := strconv.ParseInt(q.Get("task_id"), 10, 64)
switch {
case companyID > 0:
assets, err = as.QueryByCompany(companyID, typ, limit, offset)
if err == nil {
total, err = as.CountByCompany(companyID, typ)
total, err = as.CountByCompany(companyID, typ)
if err == nil && offset < total {
assets, err = as.QueryByCompany(companyID, typ, limit, offset)
}
case taskID > 0:
assets, err = as.QueryByTask(taskID, typ, limit, offset)
if err == nil {
total, err = as.CountByTask(taskID, typ)
total, err = as.CountByTask(taskID, typ)
if err == nil && offset < total {
assets, err = as.QueryByTask(taskID, typ, limit, offset)
}
default:
if typ == "" {
typ = "root_domain"
}
assets, err = as.QueryByType(typ, limit, offset)
if err == nil {
total, _ = as.CountByType(typ)
total, err = as.CountByType(typ)
if err == nil && offset < total {
assets, err = as.QueryByType(typ, limit, offset)
}
}
}
if err != nil {
writeErr(w, 400, err.Error())
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, 200, map[string]any{
+265
View File
@@ -0,0 +1,265 @@
package server
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestCompanyScopeInputsAcceptStructuredAndLegacyRules(t *testing.T) {
var inputs companyScopeInputs
if err := json.Unmarshal([]byte(`[
"example.com",
{"kind":"icp","value":"京 ICP备 123号"},
{"kind":"keyword","value":"Acme Security"}
]`), &inputs); err != nil {
t.Fatal(err)
}
if len(inputs) != 3 {
t.Fatalf("input count=%d want 3", len(inputs))
}
if inputs[0].Kind != "" || inputs[0].Value != "example.com" {
t.Fatalf("legacy rule changed: %+v", inputs[0])
}
if inputs[1].Kind != "icp" || inputs[2].Kind != "keyword" {
t.Fatalf("structured rules changed: %+v", inputs)
}
}
func TestCreateCompanyRejectsNormalizedDuplicateWithoutChangingScope(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("HTTP Strict Company %d", stamp)
companyID, _, _, _, _, err := m.pg.Companies().CreateCompanyWithScope(name, "", []db.ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("existing-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
s := &Server{m: m}
body, err := json.Marshal(map[string]any{
"name": " " + name + " ",
"scope": []map[string]string{{"kind": "domain", "value": fmt.Sprintf("new-%d.example", stamp)}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusConflict {
t.Fatalf("duplicate create status=%d body=%s", rec.Code, rec.Body.String())
}
scope, err := m.pg.Companies().GetScope(companyID)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("existing-%d.example", stamp) {
t.Fatalf("duplicate HTTP create changed existing scope: %+v", scope)
}
}
func TestDeleteCompanyRefreshesLiveTaskCompanyIDs(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete Company DTO %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("company deletion dto", "goal", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", bytes.NewBufferString(`{}`))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("delete company status=%d body=%s", rec.Code, rec.Body.String())
}
if dto := taskDTO(task, "created"); len(dto.CompanyIDs) != 0 {
t.Fatalf("live task DTO retained deleted company: %+v", dto.CompanyIDs)
}
persisted, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if persisted == nil || len(persisted.CompanyIDs) != 0 {
t.Fatalf("persisted task retained deleted company: %+v", persisted)
}
}
func TestCompanyScopeHTTPErrorClassificationAndBounds(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
t.Run("missing company is 404", func(t *testing.T) {
body := bytes.NewBufferString(`{"scope":[]}`)
req := httptest.NewRequest(http.MethodPost, "/api/companies/999999999/scope", body)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized Unicode scope is 400", func(t *testing.T) {
body, err := json.Marshal(map[string]any{
"name": "Oversized Scope",
"scope": []map[string]string{{
"kind": "keyword", "value": strings.Repeat("界", db.MaxCompanyScopeRawRunes+1),
}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("too many scope rules is 400", func(t *testing.T) {
scope := make([]map[string]string, db.MaxCompanyScopeRules+1)
for i := range scope {
scope[i] = map[string]string{"kind": "keyword", "value": fmt.Sprintf("keyword-%d", i)}
}
body, err := json.Marshal(map[string]any{"name": "Too Many Scope Rules", "scope": scope})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized request body is 413", func(t *testing.T) {
body := `{"name":"` + strings.Repeat("x", maxCompanyMutationBodyBytes+1) + `"}`
req := httptest.NewRequest(http.MethodPost, "/api/companies", strings.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
}
func TestListAssetsClassifiesValidationAndDatabaseErrors(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodGet, "/api/assets?dsl=(", nil)
rec := httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("invalid DSL status=%d body=%s", rec.Code, rec.Body.String())
}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req = httptest.NewRequest(http.MethodGet, "/api/assets?type=root_domain", nil)
rec = httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("database failure status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestDeleteCompanyRejectsBadJSONAndReportsMissing(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete JSON Company %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
for _, body := range []string{`{"delete_assets":`, `{}` + `{}`} {
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", strings.NewReader(body))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("body=%q status=%d response=%s", body, rec.Code, rec.Body.String())
}
company, err := m.pg.Companies().GetCompany(companyID)
if err != nil || company == nil {
t.Fatalf("bad JSON deleted company: company=%v err=%v", company, err)
}
}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/999999999", nil)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("missing delete status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestCompanyScopeSystemFailureIsHTTP500(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies/1/scope", bytes.NewBufferString(`{"scope":[]}`))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
}
+58 -12
View File
@@ -2,6 +2,7 @@ package server
import (
"context"
"errors"
"fmt"
"log"
"math"
@@ -10,12 +11,33 @@ import (
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/intercept"
)
const (
maxConversationRequestBytes = 64 << 10
maxConversationAgentKeyRunes = 120
maxConversationTitleRunes = 200
)
func decodeConversationRequest(w http.ResponseWriter, r *http.Request, value any) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxConversationRequestBytes)
if err := decode(r, value); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, err.Error())
}
return false
}
return true
}
// ---------- conversations (chat page) ----------
//
// A conversation is a ChatGPT-style thread bound to an agent key, independent of
@@ -46,8 +68,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id"`
}
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
if !decodeConversationRequest(w, r, &req) {
return
}
req.AgentKey = strings.TrimSpace(req.AgentKey)
@@ -55,6 +76,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "agent_key 不能为空")
return
}
if utf8.RuneCountInString(req.AgentKey) > maxConversationAgentKeyRunes {
writeErr(w, 400, fmt.Sprintf("agent_key 最多 %d 个字符", maxConversationAgentKeyRunes))
return
}
a, err := pg.GetAgentByKey(req.AgentKey)
if err != nil {
writeErr(w, 500, err.Error())
@@ -74,6 +99,10 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
if title == "" {
title = "新对话"
}
if utf8.RuneCountInString(title) > maxConversationTitleRunes {
writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes))
return
}
c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID)
if err != nil {
writeErr(w, 500, err.Error())
@@ -90,8 +119,7 @@ func (s *Server) pgUpdateConversation(w http.ResponseWriter, r *http.Request) {
var req struct {
LLMProfileID *int64 `json:"llm_profile_id"` // null clears the override
}
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
if !decodeConversationRequest(w, r, &req) {
return
}
if req.LLMProfileID != nil {
@@ -135,21 +163,39 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
if !ok {
return
}
var req struct{ Title string }
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
var req struct {
Title *string `json:"title"`
Pinned *bool `json:"pinned"`
}
if !decodeConversationRequest(w, r, &req) {
return
}
title := strings.TrimSpace(req.Title)
if title == "" {
writeErr(w, 400, "标题不能为空")
if req.Title == nil && req.Pinned == nil {
writeErr(w, 400, "至少需要提供 title 或 pinned")
return
}
if err := pg.RenameConversation(c.ID, title); err != nil {
if req.Title != nil {
title := strings.TrimSpace(*req.Title)
if title == "" {
writeErr(w, 400, "标题不能为空")
return
}
if utf8.RuneCountInString(title) > maxConversationTitleRunes {
writeErr(w, 400, fmt.Sprintf("标题最多 %d 个字符", maxConversationTitleRunes))
return
}
req.Title = &title
}
updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned})
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"ok": true})
if updated == nil {
writeErr(w, 404, "conversation not found")
return
}
writeJSON(w, 200, updated)
}
func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) {
+9 -6
View File
@@ -51,6 +51,7 @@ type TaskDTO struct {
LLMFailoverState string `json:"llm_failover_state"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []string `json:"source_task_ids"`
CompanyIDs []int64 `json:"company_ids"`
}
// TokenTotalDTO is a whole-task (all agents) token aggregate.
@@ -71,9 +72,10 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
}
func taskDTO(t *Task, status string) TaskDTO {
lifecycle := t.lifecycleSnapshot()
llmState := t.llmStateSnapshot()
sourceIDs := make([]string, 0, len(t.SourceTaskIDs))
for _, id := range t.SourceTaskIDs {
sourceIDs := make([]string, 0, len(lifecycle.SourceTaskIDs))
for _, id := range lifecycle.SourceTaskIDs {
sourceIDs = append(sourceIDs, i64s(id))
}
profileIDs := append(make([]int64, 0, len(llmState.ProfileIDs)), llmState.ProfileIDs...)
@@ -85,16 +87,17 @@ func taskDTO(t *Task, status string) TaskDTO {
Status: status,
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
CreatedUnix: t.CreatedAt,
CompletedAt: completedRFC(t.CompletedAt),
CompletedUnix: t.CompletedAt,
Paused: t.Paused,
Queued: t.Queued,
CompletedAt: completedRFC(lifecycle.CompletedAt),
CompletedUnix: lifecycle.CompletedAt,
Paused: lifecycle.Paused,
Queued: lifecycle.Queued,
LLMProfileID: llmState.ProfileID,
LLMProfileIDs: profileIDs,
ActiveLLMProfileID: llmState.ActiveID,
LLMFailoverState: llmState.FailoverState,
LLMFailoverReason: llmState.FailoverReason,
SourceTaskIDs: sourceIDs,
CompanyIDs: lifecycle.CompanyIDs,
}
}
+129 -40
View File
@@ -71,8 +71,11 @@ func preview(s string, n int) string {
const (
modelErrorRetries = 2 // model_error 收场后额外重试的次数
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
workControlWaitTimeout = 30 * time.Second
)
var errWorkControlConflict = errors.New("work control conflict")
// retryableWorkerModelError excludes errors already handled by the task router.
// In particular, a quota error after partial streaming advances the task cursor
// for the next LLM call but must not replay this whole intent on the backup.
@@ -101,10 +104,7 @@ type Engine struct {
llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
deleting sync.Map // taskID -> bool, delete barrier (no new task-owned writes)
// deleteWasPaused remembers whether BeginDelete temporarily paused an already
// running task. AbortDelete uses it to restore the exact pre-delete state.
deleteWasPaused sync.Map // taskID -> bool
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
// deleteMu makes installing the delete barrier atomic with registering a new
// task operation. Once BeginDelete returns, every admitted writer is reflected
@@ -161,7 +161,7 @@ type taskRuntime struct {
type workExecution struct {
cancel context.CancelCauseFunc
done chan struct{}
done chan error
action string // user action: pause | cancel
}
@@ -181,39 +181,33 @@ func (e *Engine) Pause(taskID string, cause error) {
}
// BeginDelete installs an execution barrier before task data/files are removed.
// The temporary pause is not a user pause: AbortDelete restores the exact state
// observed here when cleanup fails.
// The temporary pause is not a user pause. The server serializes this transition
// with lifecycle admission and tells AbortDelete whether the persisted task is
// paused/queued if cleanup fails.
func (e *Engine) BeginDelete(taskID string) bool {
e.deleteMu.Lock()
if _, loaded := e.deleting.LoadOrStore(taskID, true); loaded {
e.deleteMu.Unlock()
return false
}
wasPaused := e.IsPaused(taskID)
e.deleteWasPaused.Store(taskID, wasPaused)
e.paused.Store(taskID, true)
e.deleteMu.Unlock()
e.cancelExec(taskID, agent.AbortTaskDeleted)
return true
}
func (e *Engine) AbortDelete(taskID string) {
func (e *Engine) AbortDelete(taskID string, keepPaused bool) {
e.deleteMu.Lock()
if !e.IsDeleting(taskID) {
e.deleteMu.Unlock()
return
}
wasPaused := false
if v, ok := e.deleteWasPaused.Load(taskID); ok {
wasPaused, _ = v.(bool)
}
e.deleteWasPaused.Delete(taskID)
e.deleting.Delete(taskID)
if !wasPaused {
if !keepPaused {
e.paused.Delete(taskID)
}
e.deleteMu.Unlock()
if !wasPaused && e.m != nil {
if !keepPaused && e.m != nil {
if t, ok := e.m.Task(taskID); ok {
t.Notify()
}
@@ -291,8 +285,6 @@ func (e *Engine) StopTask(taskID string) {
e.stamped.Delete(taskID)
e.inflight.Delete(taskID)
e.coordStarted.Delete(taskID)
e.deleteWasPaused.Delete(taskID)
e.deleteMu.Lock()
e.deleting.Delete(taskID)
e.deleteMu.Unlock()
@@ -312,6 +304,17 @@ func (e *Engine) cancelExec(taskID string, cause error) {
// Resume un-pauses a task and nudges a fresh planning round. The next exec under
// it gets a fresh (uncancelled) context.
func (e *Engine) Resume(t *Task) {
// BeginDelete owns the pause barrier once deletion starts. A concurrent
// resume must never clear it and let a planner/worker re-enter while cleanup
// is waiting for task operations to drain.
if t == nil {
return
}
e.deleteMu.RLock()
defer e.deleteMu.RUnlock()
if e.IsDeleting(t.ID) {
return
}
e.paused.Delete(t.ID)
t.Notify()
}
@@ -393,14 +396,14 @@ func NewEngine(m *Manager) *Engine {
// registerWork records the cancel for the work currently running intentID.
func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) {
e.workMu.Lock()
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan struct{})}
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan error, 1)}
e.workMu.Unlock()
}
// detachWork removes the live control handle once Execute has returned. complete
// must be called after the final intent state write so a waiting cancel handler can
// safely delete the worker's blackboard output without racing a late write.
func (e *Engine) detachWork(intentID int64) (action string, complete func()) {
func (e *Engine) detachWork(intentID int64) (action string, complete func(error)) {
e.workMu.Lock()
run := e.work[intentID]
if run != nil {
@@ -413,27 +416,30 @@ func (e *Engine) detachWork(intentID int64) (action string, complete func()) {
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
e.steerMu.Unlock()
if run == nil {
return action, func() {}
return action, func(error) {}
}
return action, func() { close(run.done) }
return action, func(err error) { run.done <- err }
}
// ControlWork requests a user-visible pause or cancellation and waits until the
// worker has fully stopped writing. Cancellation cleanup is performed by the API
// handler after this returns; pause state is committed by runWorkerStep itself.
func (e *Engine) ControlWork(intentID int64, action string) error {
func (e *Engine) ControlWork(ctx context.Context, intentID int64, action string) error {
if action != "pause" && action != "cancel" {
return fmt.Errorf("unsupported work action %q", action)
}
if ctx == nil {
ctx = context.Background()
}
e.workMu.Lock()
run := e.work[intentID]
if run == nil {
e.workMu.Unlock()
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
return fmt.Errorf("%w: 意图 %d 当前没有运行中的 work(可能已结束或未被领取)", errWorkControlConflict, intentID)
}
if run.action != "" && run.action != action {
if run.action != "" {
e.workMu.Unlock()
return fmt.Errorf("意图 %d 正在执行 %s 操作", intentID, run.action)
return fmt.Errorf("%w: 意图 %d 正在执行 %s 操作", errWorkControlConflict, intentID, run.action)
}
run.action = action
done := run.done
@@ -443,7 +449,41 @@ func (e *Engine) ControlWork(intentID int64, action string) error {
}
run.cancel(cause)
e.workMu.Unlock()
<-done
timer := time.NewTimer(workControlWaitTimeout)
defer timer.Stop()
select {
case err := <-done:
return err
case <-ctx.Done():
e.releaseWorkControl(intentID, run, action)
return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, ctx.Err())
case <-timer.C:
e.releaseWorkControl(intentID, run, action)
return fmt.Errorf("等待意图 %d %s 收尾: %w", intentID, action, context.DeadlineExceeded)
}
}
// releaseWorkControl drops only this caller's reservation after its wait is
// cancelled. The work context stays cancelled; runWorkerStep recognizes the
// named cancellation cause and settles the intent into the recoverable paused
// state even if the HTTP caller has gone away.
func (e *Engine) releaseWorkControl(intentID int64, run *workExecution, action string) {
e.workMu.Lock()
if current := e.work[intentID]; current == run && current.action == action {
current.action = ""
}
e.workMu.Unlock()
}
func transitionIntentState(store *db.ExplorationStore, intentID int64, expected, state string) error {
changed, err := store.CompareAndSetIntentState(intentID, expected, state)
if err != nil {
return err
}
if !changed {
return fmt.Errorf("%w: 意图 %d 不再是 %s 状态", db.ErrIntentStateConflict, intentID, expected)
}
return nil
}
@@ -744,7 +784,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
// terminal task (goals all met → done, or failed): the run is over. A
// resume/nudge — e.g. auto-resume of the active task on restart — must NOT
// re-plan (it would burn an LLM round and re-confirm a settled result).
if isTerminalStatus(t.Status) {
if isTerminalStatus(t.lifecycleSnapshot().Status) {
return
}
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
@@ -869,7 +909,9 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
emit := func(r db.Activity) { e.emitActivity(t, r) }
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
if ectx.Err() != nil || e.IsDeleting(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "open")
if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 领取后回退失败: %v", name, t.ID, intent.ID, err)
}
return true
}
// per-work child context so the planner's kill_work can stop just this work.
@@ -909,18 +951,41 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
// unregister would always be true (unregister cancels it) → every completed
// work would be wrongly marked stopped.
workCause := context.Cause(workCtx)
killed := workCtx.Err() != nil && ectx.Err() == nil
action, completeWork := e.detachWork(intent.ID)
defer completeWork()
// A caller may stop waiting and release its in-memory reservation before the
// agent honors cancellation. The named context cause remains authoritative and
// still settles the stopped run into a recoverable state.
if action == "" {
switch {
case errors.Is(workCause, agent.AbortWorkPausedByUser):
action = "pause"
case errors.Is(workCause, agent.AbortWorkCancelledByUser):
action = "cancel"
}
}
var controlErr error
defer func() { completeWork(controlErr) }()
if action == "pause" {
_ = t.Store.SetIntentState(intent.ID, "paused")
controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused")
if controlErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 暂停状态落库失败: %v", name, t.ID, intent.ID, controlErr)
return true
}
log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID)
e.touch(t.ID)
return true
}
if action == "cancel" {
// Leave the node in running until the waiting API handler deletes the intent
// and its direct outputs transactionally. No task Notify is emitted here.
// Park the stopped run in paused before handing cleanup to the API. If the
// request disconnects after cancellation, the intent remains recoverable and
// a later cancel can finish cleanup instead of leaving a phantom running row.
controlErr = transitionIntentState(t.Store, intent.ID, "running", "paused")
if controlErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 取消栅栏落库失败: %v", name, t.ID, intent.ID, controlErr)
return true
}
log.Printf("[worker %s] task %s 意图 #%d 已停止,等待取消清理", name, t.ID, intent.ID)
e.touch(t.ID)
return true
@@ -928,14 +993,18 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// if a pause cancelled this run mid-flight, return the intent to the frontier
// so it is re-claimed on resume — the worker will resume the prior LLM
// conversation from its transcript instead of restarting from scratch.
if ectx.Err() != nil && e.IsPaused(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "open")
if ectx.Err() != nil && taskExecutionPaused(context.Cause(ectx)) {
if err := transitionIntentState(t.Store, intent.ID, "running", "open"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 任务暂停回退失败: %v", name, t.ID, intent.ID, err)
}
return true
}
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
if ectx.Err() != nil && e.isSettling(t.ID) {
_ = t.Store.SetIntentState(intent.ID, "exhausted")
if err := transitionIntentState(t.Store, intent.ID, "running", "exhausted"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 超时收尾状态落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote)
e.touch(t.ID)
return true
@@ -943,14 +1012,18 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
_ = t.Store.SetIntentState(intent.ID, "stopped")
if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d 终态停止落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
e.touch(t.ID)
return true
}
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
if killed {
_ = t.Store.SetIntentState(intent.ID, "stopped")
if err := transitionIntentState(t.Store, intent.ID, "running", "stopped"); err != nil {
log.Printf("[worker %s] task %s 意图 #%d planner 停止落库失败: %v", name, t.ID, intent.ID, err)
}
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
e.touch(t.ID)
t.Notify()
@@ -975,7 +1048,9 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
if state == "blocked" && isTaskLLMChainExhausted(err) {
_ = t.Store.SetIntentBlockedReason(intent.ID, db.IntentBlockedLLMQuota)
} else {
_ = t.Store.SetIntentState(intent.ID, state)
if stateErr := transitionIntentState(t.Store, intent.ID, "running", state); stateErr != nil {
log.Printf("[worker %s] task %s 意图 #%d 终态 %s 落库失败: %v", name, t.ID, intent.ID, state, stateErr)
}
}
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote)
e.touch(t.ID)
@@ -983,6 +1058,20 @@ func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker
return true
}
func taskExecutionPaused(cause error) bool {
var abort *agent.AbortCause
if !errors.As(cause, &abort) {
return false
}
switch abort.Code {
case "paused_by_user", "paused_by_orchestrator", "paused_on_reload", "paused_race_guard",
"queued_for_admission", "llm_unavailable_queued", "task_deleted":
return true
default:
return false
}
}
func sleepCtx(ctx context.Context, d time.Duration) (done bool) {
select {
case <-ctx.Done():
+111 -3
View File
@@ -2,7 +2,9 @@ package server
import (
"context"
"errors"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
)
@@ -33,7 +35,7 @@ func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) {
t.Fatal("kill_work cancelled the entire task context")
}
_, complete := e.detachWork(42)
complete()
complete(nil)
}
func TestControlWorkCarriesUserActionCause(t *testing.T) {
@@ -49,7 +51,7 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
done := make(chan error, 1)
go func() { done <- e.ControlWork(42, tc.action) }()
go func() { done <- e.ControlWork(context.Background(), 42, tc.action) }()
<-ctx.Done()
if code, _, _, _ := agent.AbortReason(ctx); code != tc.code {
t.Fatalf("code=%q, want %q", code, tc.code)
@@ -58,7 +60,7 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
if action != tc.action {
t.Fatalf("action=%q, want %q", action, tc.action)
}
complete()
complete(nil)
if err := <-done; err != nil {
t.Fatal(err)
}
@@ -66,6 +68,61 @@ func TestControlWorkCarriesUserActionCause(t *testing.T) {
}
}
func TestControlWorkReturnsSettlementError(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
done := make(chan error, 1)
go func() { done <- e.ControlWork(context.Background(), 42, "pause") }()
<-ctx.Done()
_, complete := e.detachWork(42)
want := errors.New("pause persistence failed")
complete(want)
if err := <-done; !errors.Is(err, want) {
t.Fatalf("ControlWork error=%v, want %v", err, want)
}
}
func TestControlWorkWaitHonorsContextAndReleasesReservation(t *testing.T) {
e := NewEngine(nil)
_, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
waitCtx, stop := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer stop()
if err := e.ControlWork(waitCtx, 42, "pause"); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("ControlWork error=%v, want deadline exceeded", err)
}
e.workMu.Lock()
run := e.work[42]
if run == nil || run.action != "" {
t.Fatalf("timed-out control left reservation: %+v", run)
}
e.workMu.Unlock()
_, complete := e.detachWork(42)
complete(nil)
}
func TestControlWorkRejectsConcurrentController(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancelCause(context.Background())
e.registerWork(42, cancel)
first := make(chan error, 1)
go func() { first <- e.ControlWork(context.Background(), 42, "pause") }()
<-ctx.Done()
if err := e.ControlWork(context.Background(), 42, "cancel"); !errors.Is(err, errWorkControlConflict) {
t.Fatalf("second control error=%v, want conflict", err)
}
action, complete := e.detachWork(42)
if action != "pause" {
t.Fatalf("winning action=%q, want pause", action)
}
complete(nil)
if err := <-first; err != nil {
t.Fatal(err)
}
}
func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) {
for _, tc := range []struct {
cause *agent.AbortCause
@@ -107,3 +164,54 @@ func TestResumeYieldsFreshContext(t *testing.T) {
t.Fatal("resume did not create a fresh execution context")
}
}
func TestTaskPauseCauseSurvivesImmediateResume(t *testing.T) {
e := NewEngine(nil)
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
old := e.execContextFor(context.Background(), task.ID)
e.Pause(task.ID, agent.AbortPausedByUser)
e.Resume(task)
if e.IsPaused(task.ID) {
t.Fatal("resume did not clear the task pause flag")
}
if old.Err() == nil || !taskExecutionPaused(context.Cause(old)) {
t.Fatalf("old execution cause=%v was not retained as a task pause", context.Cause(old))
}
if taskExecutionPaused(agent.AbortGoalMet) {
t.Fatal("goal completion was misclassified as a task pause")
}
}
func TestDeleteCancellationReturnsRunningIntentToFrontier(t *testing.T) {
if !taskExecutionPaused(agent.AbortTaskDeleted) {
t.Fatal("delete cancellation must settle a running intent as open so AbortDelete can resume it")
}
if taskExecutionPaused(agent.AbortGoalMet) {
t.Fatal("terminal task cancellation must not reopen completed work")
}
}
func TestResumeDoesNotClearDeleteBarrier(t *testing.T) {
e := NewEngine(nil)
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
e.paused.Store(task.ID, true)
e.deleting.Store(task.ID, true)
e.Resume(task)
if !e.IsPaused(task.ID) {
t.Fatal("resume cleared the pause barrier owned by task deletion")
}
}
func TestBootstrapAdmissionClearsQueuedPauseBarrier(t *testing.T) {
e := NewEngine(nil)
ctx, cancel := context.WithCancel(context.Background())
cancel()
task := &Task{ID: "t1", notify: make(chan struct{}, 1)}
s := &Server{engine: e, ctx: ctx}
e.Pause(task.ID, agent.AbortPausedByUser)
s.startAdmittedTask(task, "bootstrap")
if e.IsPaused(task.ID) {
t.Fatal("bootstrap admission left the queued task behind the pause barrier")
}
}
+16 -3
View File
@@ -71,13 +71,26 @@ func (e *Engine) taskDeadline(t *Task) int64 {
if v, ok := e.deadline.Load(t.ID); ok {
return v.(int64)
}
if t.DeadlineAt > 0 {
e.deadline.Store(t.ID, t.DeadlineAt)
return t.DeadlineAt
deadlineAt := t.lifecycleSnapshot().DeadlineAt
if deadlineAt > 0 {
e.deadline.Store(t.ID, deadlineAt)
return deadlineAt
}
return 0
}
// resetTimeoutRevival clears only the per-run timeout state after PostgreSQL has
// atomically committed timeout -> running and reset first_run_at/deadline_at. The
// configured TimeoutSeconds remains on Task, so the next real Planner/Worker run
// stamps a fresh full budget. coordStarted is reset because the coordinator that
// produced the timeout has already completed (or is in its final return path).
func (e *Engine) resetTimeoutRevival(taskID string) {
e.settling.Delete(taskID)
e.deadline.Delete(taskID)
e.stamped.Delete(taskID)
e.coordStarted.Delete(taskID)
}
// stampFirstRun records first_run_at + deadline_at on the FIRST real run (LLM ready)
// of a timeout task, once per process. No-op when the task has no timeout.
func (e *Engine) stampFirstRun(t *Task) {
+166
View File
@@ -0,0 +1,166 @@
package server
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"unicode/utf8"
"github.com/Autumn-27/artex/db"
)
const maxFindingFollowUpRunes = 4000
func findingPaginationParam(raw string, fallback, upperBound int) int {
value := atoiDefault(raw, fallback)
if value <= 0 {
value = fallback
}
if upperBound > 0 && value > upperBound {
value = upperBound
}
return value
}
// findingGroups serves the task-group page used by the global findings view.
// Findings inside a group continue to use the existing flat findings endpoint,
// preserving dashboard/export clients and giving every expanded group its own
// independent page cursor.
func (s *Server) findingGroups(w http.ResponseWriter, r *http.Request) {
q := r.URL.Query()
page := findingPaginationParam(q.Get("page"), 1, 0)
limit := findingPaginationParam(q.Get("limit"), 10, 100)
filter := db.FindingFilter{
Severity: normFilter(q.Get("severity")),
Status: normFilter(q.Get("status")),
VulnClass: normFilter(q.Get("vulnclass")),
TaskID: normFilter(q.Get("task_id")),
Sort: q.Get("sort"),
}
groups, total, findingTotal, err := s.m.pg.ListFindingGroups(filter, page, limit)
if err != nil {
writeErr(w, 500, err.Error())
return
}
// PostgreSQL owns paused/queued persistence. The live engine adds the same
// derived running/idle view used by task DTOs, so finding groups do not lag
// behind the task list while an Agent is actively executing.
if s.engine != nil {
for i := range groups {
if groups[i].TaskID == nil {
continue
}
if task, ok := s.m.Task(i64s(*groups[i].TaskID)); ok {
groups[i].TaskStatus = s.resolvedTaskStatus(task)
}
}
}
writeJSON(w, 200, map[string]any{
"items": groups,
"total": total,
"finding_total": findingTotal,
"page": page,
"page_size": limit,
})
}
// deepenFinding creates a high-priority human Worker intent under the finding's
// original task. The DB write validates and links the live finding node; task
// admission then revives or queues the task through the shared concurrency path.
func (s *Server) deepenFinding(w http.ResponseWriter, r *http.Request) {
id := int64(atoiDefault(r.PathValue("id"), 0))
if id <= 0 {
writeErr(w, 400, "bad finding id")
return
}
r.Body = http.MaxBytesReader(w, r.Body, 32<<10)
var req struct {
Description string `json:"description"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, "bad json: "+err.Error())
}
return
}
description := strings.TrimSpace(req.Description)
switch {
case description == "":
writeErr(w, 400, "description is required")
return
case utf8.RuneCountInString(description) > maxFindingFollowUpRunes:
writeErr(w, 400, fmt.Sprintf("description must be at most %d characters", maxFindingFollowUpRunes))
return
}
finding, err := s.m.pg.GetFinding(id)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if finding == nil {
writeErr(w, 404, "finding not found")
return
}
if finding.TaskID == nil || finding.NodeID == nil {
writeErr(w, 409, "finding origin task or node is no longer available")
return
}
t, ok := s.m.Task(i64s(*finding.TaskID))
if !ok || t == nil {
writeErr(w, 409, "finding origin task is no longer available")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "task is being deleted")
return
}
defer s.engine.decInflight(t.ID)
audit := db.Activity{
Worker: "system",
Kind: "text",
Summary: "人工提交漏洞深入利用意图",
Detail: description,
}
intentID, audit, err := t.Store.AddFindingFollowUpIntent(id, *finding.NodeID, description, audit)
if errors.Is(err, db.ErrFindingOriginUnavailable) {
writeErr(w, 409, "finding origin task or node is no longer available")
return
}
if err != nil {
writeErr(w, 500, err.Error())
return
}
queued, err := s.admitTask(t, "resume")
if err != nil {
if rollbackErr := t.Store.DiscardOpenIntent(intentID); rollbackErr != nil {
err = errors.Join(err, fmt.Errorf("discard follow-up intent %d: %w", intentID, rollbackErr))
}
writeErr(w, 500, err.Error())
return
}
// The audit row committed atomically with the intent. Publish that exact row;
// emitActivity would insert a second copy and break DB/SSE correspondence.
s.engine.Broadcaster().Publish(t.ID, audit)
s.engine.touch(t.ID)
// Admission can wake a Worker immediately. Report the persisted intent state
// observed at response time instead of always claiming it is still open. The
// task-level queued flag remains separate because a queued task intentionally
// keeps its Worker intent open until it receives a concurrency slot.
state := "open"
if current, stateErr := t.Store.GetNode(intentID); stateErr == nil && current != nil {
state = current.State
}
writeJSON(w, 200, map[string]any{
"task_id": t.ID,
"intent_id": i64s(intentID),
"state": state,
"queued": queued,
})
}
+321
View File
@@ -0,0 +1,321 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestFindingPaginationParam(t *testing.T) {
for _, test := range []struct {
raw, name string
fallback, upper, want int
}{
{name: "missing", raw: "", fallback: 10, upper: 100, want: 10},
{name: "invalid", raw: "not-a-number", fallback: 10, upper: 100, want: 10},
{name: "zero", raw: "0", fallback: 10, upper: 100, want: 10},
{name: "negative", raw: "-2", fallback: 1, upper: 0, want: 1},
{name: "bounded", raw: "500", fallback: 10, upper: 100, want: 100},
{name: "valid", raw: "25", fallback: 10, upper: 100, want: 25},
} {
t.Run(test.name, func(t *testing.T) {
if got := findingPaginationParam(test.raw, test.fallback, test.upper); got != test.want {
t.Fatalf("findingPaginationParam(%q,%d,%d)=%d, want %d",
test.raw, test.fallback, test.upper, got, test.want)
}
})
}
}
func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
marker := fmt.Sprintf("__server_finding_groups_%d__", time.Now().UnixNano())
createdTaskIDs := []int64{}
defer func() {
_, _ = m.pg.Exec(`DELETE FROM findings WHERE vulnclass=$1`, marker)
for _, taskID := range createdTaskIDs {
_ = m.pg.DeleteTask(taskID)
}
}()
taskA, err := m.CreateTask("finding group A", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
createdTaskIDs = append(createdTaskIDs, mustTaskID(t, taskA.ID))
taskB, err := m.CreateTask("finding group B", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
createdTaskIDs = append(createdTaskIDs, mustTaskID(t, taskB.ID))
for i, seed := range []struct {
taskID string
severity string
}{
{taskA.ID, db.SeverityCritical},
{taskA.ID, db.SeverityHigh},
{taskB.ID, db.SeverityLow},
} {
if _, err := m.pg.AddFinding(mustTaskID(t, seed.taskID), 0, marker, "", seed.severity,
fmt.Sprintf("finding %d", i), "", "test", nil); err != nil {
t.Fatalf("AddFinding[%d]: %v", i, err)
}
}
s := &Server{m: m}
request := httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&sort=severity&page=1&limit=1", nil)
recorder := httptest.NewRecorder()
s.findingGroups(recorder, request)
if recorder.Code != http.StatusOK {
t.Fatalf("groups status=%d body=%s", recorder.Code, recorder.Body.String())
}
var page struct {
Items []db.FindingGroup `json:"items"`
Total int `json:"total"`
FindingTotal int `json:"finding_total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
}
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
t.Fatal(err)
}
if len(page.Items) != 1 || page.Total != 2 || page.FindingTotal != 3 || page.Page != 1 || page.PageSize != 1 {
t.Fatalf("unexpected grouped page: %+v", page)
}
if page.Items[0].TaskID == nil || *page.Items[0].TaskID != mustTaskID(t, taskA.ID) ||
page.Items[0].Count != 2 || page.Items[0].Critical != 1 || page.Items[0].High != 1 {
t.Fatalf("unexpected first group: %+v", page.Items[0])
}
request = httptest.NewRequest(http.MethodGet,
"/api/exploration/findings/groups?vulnclass="+marker+"&page=-4&limit=0", nil)
recorder = httptest.NewRecorder()
s.findingGroups(recorder, request)
if recorder.Code != http.StatusOK {
t.Fatalf("normalized groups status=%d body=%s", recorder.Code, recorder.Body.String())
}
page = struct {
Items []db.FindingGroup `json:"items"`
Total int `json:"total"`
FindingTotal int `json:"finding_total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
}{}
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
t.Fatal(err)
}
if page.Page != 1 || page.PageSize != 10 || len(page.Items) != 2 {
t.Fatalf("pagination was not normalized: %+v", page)
}
}
func TestDeepenFindingCreatesAuditedIntentAndRevivesTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
task, err := m.CreateTask("deepen finding", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
store := task.Store
findingNodeID, err := store.AddNode(db.KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer m.pg.DeleteTask(taskID)
findingID, err := m.pg.AddFinding(taskID, findingNodeID, "test", "source", db.SeverityHigh, "source", "", "worker", nil)
if err != nil {
t.Fatal(err)
}
defer m.pg.DeleteFinding(findingID)
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel() // admission is exercised without starting background planner/worker work
s := &Server{m: m, engine: NewEngine(m), ctx: ctx}
live, unsubscribe := s.engine.Broadcaster().Subscribe(task.ID)
defer unsubscribe()
body := bytes.NewBufferString(`{"description":"验证完整利用链并保留可复现证据"}`)
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", body)
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("deepen status=%d body=%s", rec.Code, rec.Body.String())
}
var response struct {
TaskID string `json:"task_id"`
IntentID string `json:"intent_id"`
State string `json:"state"`
Queued bool `json:"queued"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &response); err != nil {
t.Fatal(err)
}
intentID, err := strconv.ParseInt(response.IntentID, 10, 64)
if err != nil || intentID <= 0 || response.TaskID != task.ID || response.State != "open" || response.Queued {
t.Fatalf("unexpected response: %+v", response)
}
intent, err := store.GetNode(intentID)
if err != nil || intent == nil {
t.Fatalf("intent=%+v err=%v", intent, err)
}
if intent.Priority != 10 || intent.Origin != "human" || intent.State != "open" {
t.Fatalf("intent metadata: %+v", intent)
}
if task.Status != "running" || task.Paused || task.Queued {
t.Fatalf("task was not revived through admission: %+v", task)
}
activity, _, err := store.ActivityList(&intentID, 0, 50)
if err != nil {
t.Fatal(err)
}
var persistedAudit db.Activity
auditCount := 0
for _, item := range activity {
if item.Worker == "system" && strings.Contains(item.Summary, "人工提交漏洞深入利用意图") {
persistedAudit = item
auditCount++
}
}
if auditCount != 1 || persistedAudit.ID <= 0 {
t.Fatalf("atomic system audit count=%d activity=%+v", auditCount, activity)
}
select {
case broadcast := <-live:
if broadcast.ID != persistedAudit.ID || broadcast.NodeID == nil || *broadcast.NodeID != intentID ||
broadcast.CreatedAt.IsZero() {
t.Fatalf("broadcast must reuse committed activity: broadcast=%+v persisted=%+v", broadcast, persistedAudit)
}
default:
t.Fatal("committed follow-up audit was not broadcast")
}
select {
case duplicate := <-live:
t.Fatalf("follow-up audit was broadcast more than once: %+v", duplicate)
default:
}
// The retained finding can no longer be deepened after its task and node are
// deleted; this must be a conflict rather than silently creating orphan work.
if err := m.pg.DeleteTask(taskID); err != nil {
t.Fatal(err)
}
req = httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", bytes.NewBufferString(`{"description":"retry"}`))
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec = httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusConflict {
t.Fatalf("deleted origin status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestDeepenFindingValidatesDescription(t *testing.T) {
s := &Server{}
tooLong, _ := json.Marshal(map[string]string{"description": strings.Repeat("x", maxFindingFollowUpRunes+1)})
for _, body := range []string{`{}`, `{"description":" "}`, string(tooLong)} {
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", bytes.NewBufferString(body))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("body=%s status=%d response=%s", body, rec.Code, rec.Body.String())
}
}
oversized := `{"description":"` + strings.Repeat("x", (32<<10)+1) + `"}`
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen", strings.NewReader(oversized))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized status=%d response=%s", rec.Code, rec.Body.String())
}
}
func TestDeepenAdmissionFailureDiscardsFollowUpIntent(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("deepen rollback", "leave no orphan intent", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer func() {
_, _ = m.pg.Exec(`UPDATE tasks SET deleted_at=NULL WHERE id=$1`, taskID)
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
findingNodeID, err := task.Store.AddNode(db.KindFinding, map[string]any{"summary": "source"}, 5, "confirmed", "worker", nil)
if err != nil {
t.Fatal(err)
}
findingID, err := m.pg.AddFinding(taskID, findingNodeID, "test", "source", db.SeverityHigh, "source", "", "worker", nil)
if err != nil {
t.Fatal(err)
}
before, err := task.Store.ListByKind(db.KindIntent, 1000)
if err != nil {
t.Fatal(err)
}
var activitiesBefore int
if err := m.pg.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.Store.ID()).Scan(&activitiesBefore); err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
// Keep the in-memory handle and graph available while forcing the atomic
// admission UPDATE to reject this soft-deleted row.
if _, err := m.pg.Exec(`UPDATE tasks SET deleted_at=now() WHERE id=$1`, taskID); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
s := &Server{m: m, engine: NewEngine(m), ctx: ctx}
req := httptest.NewRequest(http.MethodPost, "/api/exploration/findings/1/deepen",
bytes.NewBufferString(`{"description":"this intent must be rolled back"}`))
req.SetPathValue("id", strconv.FormatInt(findingID, 10))
rec := httptest.NewRecorder()
s.deepenFinding(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
after, err := task.Store.ListByKind(db.KindIntent, 1000)
if err != nil {
t.Fatal(err)
}
if len(after) != len(before) {
t.Fatalf("failed admission left an orphan follow-up: before=%d after=%d nodes=%+v", len(before), len(after), after)
}
var activitiesAfter int
if err := m.pg.QueryRow(`SELECT COUNT(*) FROM activity WHERE exploration_id=$1`, task.Store.ID()).Scan(&activitiesAfter); err != nil {
t.Fatal(err)
}
if activitiesAfter != activitiesBefore {
t.Fatalf("failed admission left an orphan audit: before=%d after=%d", activitiesBefore, activitiesAfter)
}
}
+239 -49
View File
@@ -33,14 +33,10 @@ func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
if seedFirstIntent {
s.seedFirstIntent(t)
}
if s.queueIfAtCapacity(t) {
s.engine.decInflight(t.ID)
return
s.engine.decInflight(t.ID)
if _, err := s.admitTask(t, "bootstrap"); err != nil {
log.Printf("[concurrency] task %s 启动失败: %v", t.ID, err)
}
go func() {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}()
}
func (s *Server) startTaskEngine(t *Task) {
@@ -70,7 +66,21 @@ func (s *Server) startTaskEngine(t *Task) {
}
func (s *Server) occupiesConcurrencySlot(t *Task) bool {
return t != nil && !t.Queued && !t.Paused && !s.engine.IsPaused(t.ID) && !isTerminalStatus(t.Status)
if t == nil {
return false
}
lifecycle := t.lifecycleSnapshot()
if lifecycle.Queued || lifecycle.Paused || isTerminalStatus(lifecycle.Status) {
return false
}
// Deletion temporarily pauses the Engine but has not committed yet. Preserve
// the task's slot until PostgreSQL deletion succeeds; otherwise FIFO promotion
// during the drain window could over-admit if deletion later aborts and the
// persisted running task is restored.
if s.engine.IsDeleting(t.ID) {
return true
}
return !s.engine.IsPaused(t.ID) && (s.engine.ReadyFor(t) || s.engine.ActiveLLMCalls(t.ID) > 0)
}
func (s *Server) runningTaskCount(excludeID string) int {
@@ -83,58 +93,255 @@ func (s *Server) runningTaskCount(excludeID string) int {
return count
}
func (s *Server) queueIfAtCapacity(t *Task) bool {
// admitTask is the single admission path for new, resumed, rerun and follow-up
// work. It atomically either starts the task or appends it to the persistent FIFO
// queue. mode is bootstrap for a freshly-created task and resume otherwise.
func (s *Server) admitTask(t *Task, mode string) (queued bool, err error) {
return s.admitTaskWhen(t, mode, false)
}
// admitPausedTask is the task-control resume path. The paused precondition is
// checked under the same scheduler lock as admission so a concurrent pause,
// dequeue or FIFO promotion cannot leave database and Engine state divergent.
func (s *Server) admitPausedTask(t *Task) (queued bool, err error) {
return s.admitTaskWhen(t, "resume", true)
}
func (s *Server) admitTaskWhen(t *Task, mode string, requirePaused bool) (queued bool, err error) {
if t == nil {
return false, fmt.Errorf("task not found")
}
if mode != "bootstrap" {
mode = "resume"
}
s.concMu.Lock()
defer s.concMu.Unlock()
// Delete installs its barrier under concMu as well. Re-resolve after acquiring
// the lock so a request that captured a task pointer before successful deletion
// cannot revive that stale handle after StopTask clears its Engine maps.
current, exists := s.m.Task(t.ID)
if !exists || current != t || s.engine.IsDeleting(t.ID) {
return false, fmt.Errorf("task is being deleted")
}
if !s.engine.beginTaskOperation(t.ID) {
return false, fmt.Errorf("task is being deleted")
}
defer s.engine.decInflight(t.ID)
lifecycle := t.lifecycleSnapshot()
if requirePaused {
if isTerminalStatus(lifecycle.Status) {
return false, fmt.Errorf("终态任务不能执行继续")
}
if !lifecycle.Paused {
return false, fmt.Errorf("仅已暂停的任务可以继续")
}
mode = s.resumeAdmissionMode(t)
}
wasTerminal := isTerminalStatus(lifecycle.Status)
wasPaused := lifecycle.Paused || s.engine.IsPaused(t.ID)
wasQueued := lifecycle.Queued
engineWasPaused := s.engine.IsPaused(t.ID)
enabled, limit := s.m.ConcurrencyLimit()
if !enabled || s.runningTaskCount(t.ID) < limit {
return false
ready := s.engine.ReadyFor(t)
// Work added to a task that is already admitted only needs a wake-up. This
// matters when the configured limit was lowered below the current running
// count: an existing task must not suddenly mark itself queued while its
// planner/workers are still live. Still compare-and-commit the persisted
// status: a concurrent terminal transition must win instead of being silently
// reported as a successful follow-up admission.
if !wasTerminal && !wasPaused && !wasQueued && s.engine.Started(t.ID) && (!enabled || ready) {
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, lifecycle.Status, false, "resume", false); err != nil {
return false, err
}
s.startAdmittedTask(t, "resume")
return false, nil
}
if err := s.m.SetTaskQueued(t.ID, true); err != nil {
log.Printf("[concurrency] task %s 置排队失败: %v", t.ID, err)
return false
// Preserve the original first-run mode across repeated admissions. Legacy
// queued rows have an empty queue_mode, so infer bootstrap from their graph.
if wasQueued {
switch lifecycle.QueueMode {
case "bootstrap":
mode = "bootstrap"
case "":
mode = s.resumeAdmissionMode(t)
}
}
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text",
Summary: fmt.Sprintf("已排队:达到并发上限 %d,等待空位后自动开始", limit)})
return true
readyBacklog := enabled && s.hasReadyQueuedTask(t.ID)
atCapacity := enabled && s.runningTaskCount(t.ID) >= limit
shouldQueue := enabled && (!ready || readyBacklog || atCapacity)
// Install the execution barrier before reviving a terminal/paused task. Without
// this ordering, its already-running worker loops can claim the newly-opened
// intent in the gap between status=running and queued=true.
if shouldQueue || wasTerminal || wasPaused || wasQueued {
s.engine.Pause(t.ID, agent.Causef("queued_for_admission", "任务等待运行准入",
"任务正在等待并发队列或准入状态提交,本次执行已停止;只有获得运行槽后才会重新领取意图"))
}
status := lifecycle.Status
if wasTerminal {
status = "running"
}
if err := s.m.ApplyTaskAdmission(t.ID, lifecycle.Status, status, shouldQueue, mode, wasQueued); err != nil {
if !engineWasPaused && !wasPaused && !wasQueued {
s.engine.Resume(t)
}
return false, err
}
if lifecycle.Status == "timeout" && status == "running" {
// ApplyTaskAdmission reset this timed-out run's persisted clock. Clear the
// matching Engine gates before either parking it in FIFO or starting work;
// otherwise the old settling flag would make every worker skip forever.
s.engine.resetTimeoutRevival(t.ID)
}
if shouldQueue {
if !wasQueued {
summary := fmt.Sprintf("已排队:达到并发上限 %d,等待空位后自动开始", limit)
switch {
case !ready:
summary = "已排队:当前没有可运行的 LLM 配置,配置恢复后自动开始"
case readyBacklog:
summary = "已排队:已有更早的任务等待运行,将按 FIFO 顺序自动开始"
}
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text", Summary: summary})
}
return true, nil
}
s.startAdmittedTask(t, mode)
return false, nil
}
func (s *Server) hasReadyQueuedTask(excludeID string) bool {
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if task.ID != excludeID && lifecycle.Queued && !lifecycle.Paused && !isTerminalStatus(lifecycle.Status) && s.engine.ReadyFor(task) {
return true
}
}
return false
}
func (s *Server) startAdmittedTask(t *Task, mode string) {
if mode == "bootstrap" {
if !s.engine.beginTaskOperation(t.ID) {
return
}
// A first-run task may have kept the Engine pause barrier while waiting
// in the concurrency queue. Clear it only after operation admission, or
// startTaskEngine's first execContextFor call would return a cancelled
// context and silently strand the dequeued task.
if s.engine.IsPaused(t.ID) {
s.engine.Resume(t)
}
go func() {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}()
return
}
s.engine.Run(s.ctx, t)
s.engine.Resume(t)
// Run returns early for an already-started task. Explicitly ensure a timeout
// coordinator exists after resume; resetTimeoutRevival cleared the completed
// coordinator's marker and the next real call will stamp a fresh deadline.
s.engine.startDeadlineCoordinator(s.ctx, t)
}
func (s *Server) reconcileConcurrency() {
s.concMu.Lock()
defer s.concMu.Unlock()
queued := []*Task{}
for _, task := range s.m.List() {
if task.Queued && !isTerminalStatus(task.Status) {
queued = append(queued, task)
enabled, limit := s.m.ConcurrencyLimit()
// A task whose provider chain becomes unavailable cannot do useful work and
// must not reserve a limited running slot forever. Park it persistently so a
// later profile edit/recovery can re-enter through the same FIFO path.
if enabled {
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if lifecycle.Queued || lifecycle.Paused || s.engine.IsPaused(task.ID) || s.engine.IsDeleting(task.ID) ||
isTerminalStatus(lifecycle.Status) || s.engine.ReadyFor(task) || s.engine.ActiveLLMCalls(task.ID) > 0 {
continue
}
mode := s.resumeAdmissionMode(task)
s.engine.Pause(task.ID, agent.Causef("llm_unavailable_queued", "LLM 不可用,任务进入等待队列",
"任务当前无法解析可运行的 Planner/Worker LLM,已释放并发槽;配置恢复后按队列顺序继续"))
if err := s.m.EnqueueTask(task.ID, mode); err != nil {
s.engine.Resume(task)
log.Printf("[concurrency] task %s 因 LLM 不可用入队失败: %v", task.ID, err)
continue
}
s.engine.emitActivity(task, db.Activity{Worker: "system", Kind: "text",
Summary: "已排队:当前没有可运行的 LLM 配置,配置恢复后自动开始"})
}
}
sort.Slice(queued, func(i, j int) bool { return queued[i].CreatedAt < queued[j].CreatedAt })
enabled, limit := s.m.ConcurrencyLimit()
type queuedTask struct {
task *Task
lifecycle taskLifecycleState
}
queued := []queuedTask{}
for _, task := range s.m.List() {
lifecycle := task.lifecycleSnapshot()
if lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
queued = append(queued, queuedTask{task: task, lifecycle: lifecycle})
}
}
sort.SliceStable(queued, func(i, j int) bool {
left, right := queued[i].lifecycle.QueuedAt, queued[j].lifecycle.QueuedAt
if left == 0 {
left = queued[i].task.CreatedAt * int64(1e9)
}
if right == 0 {
right = queued[j].task.CreatedAt * int64(1e9)
}
if left != right {
return left < right
}
// Old rows may not have queued_at and task creation timestamps are only
// kept to second precision in memory. Task ids are monotonic, providing a
// deterministic oldest-first fallback for those ties.
leftID, leftErr := strconv.ParseInt(queued[i].task.ID, 10, 64)
rightID, rightErr := strconv.ParseInt(queued[j].task.ID, 10, 64)
if leftErr == nil && rightErr == nil {
return leftID < rightID
}
return queued[i].task.ID < queued[j].task.ID
})
slots := len(queued)
if enabled {
slots = limit - s.runningTaskCount("")
}
for _, task := range queued {
for _, entry := range queued {
if slots <= 0 {
break
}
task := entry.task
lifecycle := task.lifecycleSnapshot()
// Keep FIFO order, but do not consume a concurrency slot for a task
// whose explicit chain/global provider is unavailable. It will be retried
// after the user configures or resets its LLM chain.
if task.Paused || !s.engine.ReadyFor(task) {
if lifecycle.Paused || (enabled && !s.engine.ReadyFor(task)) {
continue
}
mode := lifecycle.QueueMode
if mode != "bootstrap" && mode != "resume" {
mode = s.resumeAdmissionMode(task)
}
if mode != "bootstrap" {
mode = "resume"
}
if !s.engine.beginTaskOperation(task.ID) {
continue
}
if err := s.m.SetTaskQueued(task.ID, false); err != nil {
if err := s.m.ApplyTaskAdmission(task.ID, lifecycle.Status, lifecycle.Status, false, mode, false); err != nil {
s.engine.decInflight(task.ID)
continue
}
go func(t *Task) {
defer s.engine.decInflight(t.ID)
s.startTaskEngine(t)
}(task)
s.startAdmittedTask(task, mode)
s.engine.decInflight(task.ID)
slots--
}
}
@@ -149,25 +356,8 @@ func (s *Server) reviveTask(t *Task) {
if t == nil {
return
}
if t.Queued {
return
}
// 终态 → 拉回 running(SetTaskStatus 会清 completed_at 并同步内存 handle 的 Status,
// 使 planner/worker 循环的终态门放行)。
if isTerminalStatus(t.Status) {
if err := s.m.SetTaskStatus(t.ID, "running"); err != nil {
log.Printf("[revive] task %s 置 running 失败: %v", t.ID, err)
}
}
// 确保引擎循环存活:已 started 只会 Notify;未 started(重启后未恢复的终态任务)则重起循环。
s.engine.Run(s.ctx, t)
// 暂停中 → 解除暂停(清引擎内存态 + Notify),并持久化 paused=false(存活过重启)。
if t.Paused || s.engine.IsPaused(t.ID) {
s.engine.Resume(t)
t.Paused = false
if err := s.m.SetTaskPaused(t.ID, false); err != nil {
log.Printf("[revive] task %s 解除暂停失败: %v", t.ID, err)
}
if _, err := s.admitTask(t, "resume"); err != nil {
log.Printf("[revive] task %s 恢复失败: %v", t.ID, err)
}
}
+375 -52
View File
@@ -1,6 +1,7 @@
package server
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
@@ -26,21 +27,28 @@ import (
// sharing the process-wide asset store. ID is the PG task id as a string; ExpID
// is the exploration the task owns.
type Task struct {
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
Paused bool `json:"paused"`
Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
Paused bool `json:"paused"`
Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑
// QueuedAt is an internal Unix-nanosecond ordering key. It is deliberately
// finer than CreatedAt so several tasks enqueued in the same second retain
// their real FIFO order.
QueuedAt int64 `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
LLMChainRevision int64 `json:"-"`
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态;空/其它则由运行态推导)
// 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。
TimeoutSeconds int `json:"timeout_seconds"`
@@ -50,6 +58,7 @@ type Task struct {
Store *pgdb.ExplorationStore `json:"-"`
Guard *guard.Guard `json:"-"`
notify chan struct{}
lifecycleMu sync.RWMutex
llmMu sync.RWMutex
// pendingTriggers accumulates the concrete changes (worker done / finding) that
@@ -59,10 +68,72 @@ type Task struct {
pendingTriggers []agent.TriggerEvent
}
// taskLifecycleState is an internally consistent view of the mutable task
// lifecycle and inherited-scope context. Callers must use lifecycleSnapshot and
// updateLifecycle instead of reading or writing the corresponding Task fields
// directly after the task has been published by Manager.
type taskLifecycleState struct {
Status string
Paused bool
Queued bool
QueuedAt int64
QueueMode string
CompletedAt int64
FirstRunAt int64
DeadlineAt int64
SourceTaskIDs []int64
CompanyIDs []int64
}
func (t *Task) lifecycleSnapshot() taskLifecycleState {
if t == nil {
return taskLifecycleState{}
}
t.lifecycleMu.RLock()
defer t.lifecycleMu.RUnlock()
return t.lifecycleSnapshotLocked()
}
func (t *Task) lifecycleSnapshotLocked() taskLifecycleState {
return taskLifecycleState{
Status: t.Status,
Paused: t.Paused,
Queued: t.Queued,
QueuedAt: t.QueuedAt,
QueueMode: t.QueueMode,
CompletedAt: t.CompletedAt,
FirstRunAt: t.FirstRunAt,
DeadlineAt: t.DeadlineAt,
SourceTaskIDs: append([]int64(nil), t.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), t.CompanyIDs...),
}
}
func (t *Task) updateLifecycle(update func(*taskLifecycleState)) {
if t == nil || update == nil {
return
}
t.lifecycleMu.Lock()
state := t.lifecycleSnapshotLocked()
update(&state)
t.Status = state.Status
t.Paused = state.Paused
t.Queued = state.Queued
t.QueuedAt = state.QueuedAt
t.QueueMode = state.QueueMode
t.CompletedAt = state.CompletedAt
t.FirstRunAt = state.FirstRunAt
t.DeadlineAt = state.DeadlineAt
t.SourceTaskIDs = append(t.SourceTaskIDs[:0], state.SourceTaskIDs...)
t.CompanyIDs = append(t.CompanyIDs[:0], state.CompanyIDs...)
t.lifecycleMu.Unlock()
}
type taskLLMState struct {
ProfileID *int64
ProfileIDs []int64
ActiveID *int64
ChainRevision int64
FailoverState string
FailoverReason string
}
@@ -82,19 +153,25 @@ func (t *Task) llmStateSnapshot() taskLLMState {
ProfileID: cloneInt64Ptr(t.LLMProfileID),
ProfileIDs: append(make([]int64, 0, len(t.LLMProfileIDs)), t.LLMProfileIDs...),
ActiveID: cloneInt64Ptr(t.ActiveLLMProfileID),
ChainRevision: t.LLMChainRevision,
FailoverState: t.LLMFailoverState,
FailoverReason: t.LLMFailoverReason,
}
}
func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, state, reason string) {
func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, revision int64, state, reason string) bool {
t.llmMu.Lock()
defer t.llmMu.Unlock()
if revision < t.LLMChainRevision {
return false
}
t.LLMProfileID = cloneInt64Ptr(profileID)
t.ActiveLLMProfileID = cloneInt64Ptr(activeID)
t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...)
t.LLMChainRevision = revision
t.LLMFailoverState = state
t.LLMFailoverReason = reason
t.llmMu.Unlock()
return true
}
// DeleteTaskOptions controls cleanup of data stored outside the task's own
@@ -129,11 +206,16 @@ type Manager struct {
enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP)
interceptor *intercept.Interceptor // user-configured tool-call interception rules
mu sync.RWMutex
tasks map[string]*Task
active string
trafficOn bool // 流量捕获开关(默认关;settings.traffic_capture)
llmRecOn bool // LLM 录制开关(默认关;settings.llm_record)
companyMu sync.Mutex // serializes task/company-scope commits with live handle registration
// taskStateMu preserves commit order between PostgreSQL lifecycle writes and
// their in-memory mirrors. lifecycleMu makes snapshots race-free, but without
// this outer write lock an older request could commit first and publish last.
taskStateMu sync.Mutex
mu sync.RWMutex
tasks map[string]*Task
active string
trafficOn bool // 流量捕获开关(默认关;settings.traffic_capture)
llmRecOn bool // LLM 录制开关(默认关;settings.llm_record)
// 联网搜索开关与来源(默认关;settings.web_search_*)。brave-free 需要 braveKey;tavily 需要 tavilyKey。
// webSearchProxy 是独立出口代理(http/https/socks5),与记录流量的 MITM 代理无关。
webSearchOn bool
@@ -594,15 +676,25 @@ func unixOrZero(t *time.Time) int64 {
return t.Unix()
}
func unixNanoOrZero(t *time.Time) int64 {
if t == nil {
return 0
}
return t.UnixNano()
}
func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task {
return &Task{
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued,
QueuedAt: unixNanoOrZero(pt.QueuedAt), QueueMode: pt.QueueMode,
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
LLMProfileID: pt.LLMProfileID,
LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID,
LLMChainRevision: pt.LLMChainRevision,
LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason,
SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...),
CompanyIDs: append([]int64(nil), pt.CompanyIDs...),
TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds,
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
@@ -622,6 +714,10 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time
}
func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.TaskCreateOptions) (*Task, error) {
if len(opts.CompanyIDs) > 0 {
m.companyMu.Lock()
defer m.companyMu.Unlock()
}
pt, err := m.pg.CreateTaskWithOptions(description, goal, opts)
if err != nil {
return nil, err
@@ -634,6 +730,33 @@ func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.Task
return t, nil
}
// DeleteCompanyWithAssets keeps the database cascade and live task handles in
// one manager-level critical section. This closes the gap where a task could
// commit its company scope immediately before registration and miss the
// post-delete in-memory sweep.
func (m *Manager) DeleteCompanyWithAssets(id int64, deleteAssets bool) (int64, error) {
m.companyMu.Lock()
defer m.companyMu.Unlock()
assetsDeleted, err := m.pg.Companies().DeleteCompanyWithAssets(id, deleteAssets)
if err != nil {
return 0, err
}
m.mu.Lock()
for _, task := range m.tasks {
task.updateLifecycle(func(state *taskLifecycleState) {
companyIDs := make([]int64, 0, len(state.CompanyIDs))
for _, companyID := range state.CompanyIDs {
if companyID != id {
companyIDs = append(companyIDs, companyID)
}
}
state.CompanyIDs = companyIDs
})
}
m.mu.Unlock()
return assetsDeleted, nil
}
// ReplaceTaskLLMProfiles resets a non-terminal task's ordered provider chain and
// mirrors the committed state onto the live task handle.
func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activeProfileID int64) (int64, error) {
@@ -650,7 +773,7 @@ func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activePr
}
m.mu.Lock()
if task := m.tasks[id]; task != nil {
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason)
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMChainRevision, pt.LLMFailoverState, pt.LLMFailoverReason)
}
m.mu.Unlock()
if task, ok := m.Task(id); ok {
@@ -701,26 +824,212 @@ func (m *Manager) LoadExisting() []*Task {
// SetTaskPaused persists a task's paused state.
func (m *Manager) SetTaskPaused(id string, paused bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
return m.pg.SetPaused(n, paused)
}
// SetTaskQueued persists a task's queued (concurrency-hold) state and syncs the
// in-memory handle so listTasks/占位统计 see it immediately.
func (m *Manager) SetTaskQueued(id string, queued bool) error {
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if err := m.pg.SetQueued(n, queued); err != nil {
if err := m.pg.SetPaused(n, paused); err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Queued = queued
t.updateLifecycle(func(state *taskLifecycleState) {
state.Paused = paused
})
}
m.mu.Unlock()
return nil
}
// ApplyTaskAdmission atomically commits the lifecycle fields controlled by the
// concurrency scheduler. Keeping status, paused and queue metadata in one UPDATE
// prevents a failed resume from leaving a task half-revived (for example running
// but still user-paused, or dequeued without an Engine start).
//
// preservePosition applies only when the row is already queued. A repeated
// admission keeps its FIFO timestamp; a task that was explicitly paused and is
// now re-queued receives a fresh tail position.
func (m *Manager) ApplyTaskAdmission(id, expectedStatus, status string, queued bool, mode string, preservePosition bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if queued {
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
} else {
mode = ""
}
var queuedAt, completedAt, firstRunAt, deadlineAt sql.NullTime
var committedMode string
err = m.pg.QueryRow(`UPDATE tasks
SET status=$2,
completed_at=CASE
WHEN $2 IN ('done','failed','timeout') THEN COALESCE(completed_at, now())
ELSE NULL
END,
paused=false,
queued=$3,
queued_at=CASE
WHEN NOT $3 THEN NULL
WHEN $5 AND queued THEN COALESCE(queued_at, now())
ELSE now()
END,
queue_mode=CASE
WHEN NOT $3 THEN ''
WHEN ($5 AND queued AND queue_mode='bootstrap') OR $4='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END,
first_run_at=CASE
WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL
ELSE first_run_at
END,
deadline_at=CASE
WHEN $6='timeout' AND $2 NOT IN ('done','failed','timeout') THEN NULL
ELSE deadline_at
END
WHERE id=$1 AND deleted_at IS NULL AND status=$6
RETURNING queued_at, queue_mode, completed_at, first_run_at, deadline_at`, n, status, queued, mode, preservePosition, expectedStatus).
Scan(&queuedAt, &committedMode, &completedAt, &firstRunAt, &deadlineAt)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s lifecycle changed before admission (expected status %q)", id, expectedStatus)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
state.Paused = false
state.Queued = queued
state.QueueMode = committedMode
state.QueuedAt = 0
if queuedAt.Valid {
state.QueuedAt = queuedAt.Time.UnixNano()
}
state.CompletedAt = 0
if completedAt.Valid {
state.CompletedAt = completedAt.Time.Unix()
}
state.FirstRunAt = 0
if firstRunAt.Valid {
state.FirstRunAt = firstRunAt.Time.Unix()
}
state.DeadlineAt = 0
if deadlineAt.Valid {
state.DeadlineAt = deadlineAt.Time.Unix()
}
})
}
m.mu.Unlock()
return nil
}
// ApplyTaskPause atomically removes a task from the admission queue and records
// the user pause. queue_mode is intentionally retained so resuming a never-run
// bootstrap task still performs goal decomposition, but the next enqueue receives
// a new queued_at timestamp and therefore moves to the FIFO tail.
func (m *Manager) ApplyTaskPause(id string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
var mode string
err = m.pg.QueryRow(`UPDATE tasks
SET paused=true, queued=false, queued_at=NULL
WHERE id=$1 AND deleted_at IS NULL AND paused=false
AND status NOT IN ('done','failed','timeout')
RETURNING COALESCE(queue_mode,'')`, n).Scan(&mode)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s is unavailable for pause", id)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Paused = true
state.Queued = false
state.QueuedAt = 0
state.QueueMode = mode
})
}
m.mu.Unlock()
return nil
}
// EnqueueTask persists the concurrency hold and syncs the in-memory handle.
func (m *Manager) EnqueueTask(id, mode string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if mode != "bootstrap" && mode != "resume" {
return fmt.Errorf("invalid queue mode %q", mode)
}
var queuedAt time.Time
var committedMode string
err = m.pg.QueryRow(`UPDATE tasks
SET queued=true,
queued_at=CASE WHEN queued THEN COALESCE(queued_at, now()) ELSE now() END,
queue_mode=CASE
WHEN queue_mode='bootstrap' OR $2='bootstrap' THEN 'bootstrap'
ELSE 'resume'
END
WHERE id=$1 AND deleted_at IS NULL
RETURNING queued_at, queue_mode`, n, mode).Scan(&queuedAt, &committedMode)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("task %s is unavailable for enqueue", id)
}
if err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.QueuedAt = queuedAt.UnixNano()
state.Queued = true
state.QueueMode = committedMode
})
}
m.mu.Unlock()
return nil
}
// DequeueTask removes the concurrency hold. clearMode=false is used when a user
// pauses a queued task so a later resume still knows whether bootstrap is needed.
func (m *Manager) DequeueTask(id string, clearMode bool) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
}
if err := m.pg.Dequeue(n, clearMode); err != nil {
return err
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.updateLifecycle(func(state *taskLifecycleState) {
state.Queued = false
state.QueuedAt = 0
if clearMode {
state.QueueMode = ""
}
})
}
m.mu.Unlock()
return nil
@@ -731,7 +1040,7 @@ func (m *Manager) TaskStatus(id string) string {
m.mu.RLock()
defer m.mu.RUnlock()
if t := m.tasks[id]; t != nil {
return t.Status
return t.lifecycleSnapshot().Status
}
return ""
}
@@ -739,6 +1048,8 @@ func (m *Manager) TaskStatus(id string) string {
// StampTaskFirstRun stamps first_run_at + deadline_at on the first real run (idempotent
// in DB) and mirrors deadline_at on the live handle. Returns the deadline unix (0 = 不限).
func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return 0, err
@@ -759,10 +1070,12 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
if t.FirstRunAt == 0 {
t.FirstRunAt = time.Now().Unix()
}
t.DeadlineAt = dlUnix
t.updateLifecycle(func(state *taskLifecycleState) {
if state.FirstRunAt == 0 {
state.FirstRunAt = time.Now().Unix()
}
state.DeadlineAt = dlUnix
})
}
m.mu.Unlock()
return dlUnix, nil
@@ -772,6 +1085,8 @@ func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
// (resolves the completed↔timeout race — first terminal writer wins). Reflects the
// won status on the live handle. won=false means another terminal already stuck.
func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return false, err
@@ -782,10 +1097,12 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error)
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Status = status
if t.CompletedAt == 0 {
t.CompletedAt = time.Now().Unix()
}
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
if state.CompletedAt == 0 {
state.CompletedAt = time.Now().Unix()
}
})
}
m.mu.Unlock()
return true, nil
@@ -794,6 +1111,8 @@ func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error)
// SetTaskStatus persists a task's lifecycle status (e.g. "done") and reflects it
// on the in-memory handle so the derived DTO status shows it without a reload.
func (m *Manager) SetTaskStatus(id, status string) error {
m.taskStateMu.Lock()
defer m.taskStateMu.Unlock()
n, err := strconv.ParseInt(id, 10, 64)
if err != nil {
return err
@@ -803,16 +1122,18 @@ func (m *Manager) SetTaskStatus(id, status string) error {
}
m.mu.Lock()
if t := m.tasks[id]; t != nil {
t.Status = status
// mirror the DB's completed_at stamp on the live handle so the DTO shows the
// finish time without a reload (terminal → stamp once; else clear).
if pgdb.IsTerminal(status) {
if t.CompletedAt == 0 {
t.CompletedAt = time.Now().Unix()
t.updateLifecycle(func(state *taskLifecycleState) {
state.Status = status
// Mirror the DB's completed_at stamp on the live handle so the DTO shows
// the finish time without a reload (terminal -> stamp once; else clear).
if pgdb.IsTerminal(status) {
if state.CompletedAt == 0 {
state.CompletedAt = time.Now().Unix()
}
} else {
state.CompletedAt = 0
}
} else {
t.CompletedAt = 0
}
})
}
m.mu.Unlock()
return nil
@@ -927,13 +1248,15 @@ func (m *Manager) forgetTask(id string, numericID int64) {
m.mu.Lock()
delete(m.tasks, id)
for _, task := range m.tasks {
kept := task.SourceTaskIDs[:0]
for _, sourceID := range task.SourceTaskIDs {
if sourceID != numericID {
kept = append(kept, sourceID)
task.updateLifecycle(func(state *taskLifecycleState) {
kept := make([]int64, 0, len(state.SourceTaskIDs))
for _, sourceID := range state.SourceTaskIDs {
if sourceID != numericID {
kept = append(kept, sourceID)
}
}
}
task.SourceTaskIDs = kept
state.SourceTaskIDs = kept
})
}
if m.active == id {
m.active = ""
+109
View File
@@ -0,0 +1,109 @@
package server
import (
"fmt"
"sync"
"testing"
)
func TestTaskLifecycleSnapshotCopiesContextSlices(t *testing.T) {
task := &Task{
Status: "running",
SourceTaskIDs: []int64{11, 12},
CompanyIDs: []int64{21, 22},
}
snapshot := task.lifecycleSnapshot()
snapshot.SourceTaskIDs[0] = 99
snapshot.CompanyIDs[0] = 98
fresh := task.lifecycleSnapshot()
if got := fresh.SourceTaskIDs[0]; got != 11 {
t.Fatalf("snapshot source IDs alias task storage: got %d", got)
}
if got := fresh.CompanyIDs[0]; got != 21 {
t.Fatalf("snapshot company IDs alias task storage: got %d", got)
}
sources := []int64{31, 32}
companies := []int64{41, 42}
task.updateLifecycle(func(state *taskLifecycleState) {
state.SourceTaskIDs = sources
state.CompanyIDs = companies
})
sources[0] = 97
companies[0] = 96
fresh = task.lifecycleSnapshot()
if got := fresh.SourceTaskIDs[0]; got != 31 {
t.Fatalf("lifecycle update retained caller source slice: got %d", got)
}
if got := fresh.CompanyIDs[0]; got != 41 {
t.Fatalf("lifecycle update retained caller company slice: got %d", got)
}
}
func TestTaskLifecycleSnapshotConcurrentConsistency(t *testing.T) {
task := &Task{}
writeState := func(generation int64, queued bool) {
task.updateLifecycle(func(state *taskLifecycleState) {
state.Status = map[bool]string{true: "queued", false: "paused"}[queued]
state.Queued = queued
state.Paused = !queued
state.QueuedAt = generation
state.QueueMode = map[bool]string{true: "resume", false: "bootstrap"}[queued]
state.CompletedAt = generation
state.FirstRunAt = generation
state.DeadlineAt = generation
state.SourceTaskIDs = []int64{generation, generation + 1}
state.CompanyIDs = []int64{generation, generation + 1}
})
}
writeState(1, true)
var wg sync.WaitGroup
for writer := 0; writer < 4; writer++ {
writer := writer
wg.Add(1)
go func() {
defer wg.Done()
for i := 1; i <= 1000; i++ {
generation := int64(writer*1000 + i)
writeState(generation, generation%2 == 0)
}
}()
}
errCh := make(chan error, 4)
for reader := 0; reader < 4; reader++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 1000; i++ {
state := task.lifecycleSnapshot()
if state.Queued == state.Paused {
errCh <- fmt.Errorf("queued=%v paused=%v", state.Queued, state.Paused)
return
}
wantStatus := map[bool]string{true: "queued", false: "paused"}[state.Queued]
wantMode := map[bool]string{true: "resume", false: "bootstrap"}[state.Queued]
if state.Status != wantStatus || state.QueueMode != wantMode {
errCh <- fmt.Errorf("inconsistent status/mode: %+v", state)
return
}
generation := state.QueuedAt
if state.CompletedAt != generation || state.FirstRunAt != generation || state.DeadlineAt != generation ||
len(state.SourceTaskIDs) != 2 || len(state.CompanyIDs) != 2 ||
state.SourceTaskIDs[0] != generation || state.CompanyIDs[0] != generation {
errCh <- fmt.Errorf("mixed lifecycle snapshot: %+v", state)
return
}
}
}()
}
wg.Wait()
close(errCh)
for err := range errCh {
t.Fatal(err)
}
}
+9 -6
View File
@@ -302,11 +302,13 @@ func (s *Server) toolPauseTask() actool.CoreTool {
TaskID string `json:"task_id"`
}
_ = json.Unmarshal(in, &a)
if _, ok := s.m.Task(a.TaskID); !ok {
t, ok := s.m.Task(a.TaskID)
if !ok {
return actool.Errorf("task 不存在: " + a.TaskID), nil
}
s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator)
s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask)
if _, err := s.applyTaskControlWithCause(t, "pause", agent.AbortPausedByOrchestrator); err != nil {
return actool.Errorf(err.Error()), nil
}
return actool.Text("task paused: " + a.TaskID), nil
})
}
@@ -416,10 +418,11 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
func (s *Server) deriveTaskStatus(t *Task) string {
lifecycle := t.lifecycleSnapshot()
switch {
case isTerminalStatus(t.Status):
return t.Status
case t.Paused || s.engine.IsPaused(t.ID):
case isTerminalStatus(lifecycle.Status):
return lifecycle.Status
case lifecycle.Paused || s.engine.IsPaused(t.ID):
return "paused"
case s.engine.ReadyFor(t) && s.engine.Started(t.ID):
return "running"
+151 -100
View File
@@ -4,7 +4,7 @@ import (
"archive/zip"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
@@ -238,26 +238,32 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
// reload tasks persisted on disk so the task list survives a restart, and
// restore persisted paused state (so a task paused before restart stays paused).
for _, t := range m.LoadExisting() {
lifecycle := t.lifecycleSnapshot()
// clear stale 'running' intents from a prior crash/restart (no live worker
// owns them) so they re-claim instead of spinning forever in the UI.
if n, _ := t.Store.ResetRunningIntents(); n > 0 {
log.Printf("[engine] task %s 重置 %d 个残留 running 意图为 open", t.ID, n)
}
if t.Paused {
if lifecycle.Paused {
s.engine.Pause(t.ID, agent.AbortPausedOnReload)
}
// 任务级超时:为每个未终态、带 timeout 的任务起 deadline 协调器,独立于 planner/worker
// loop——非活跃任务重启后也能在到点后被收尾(deadline 已过则立即走收尾时序)。
if !isTerminalStatus(t.Status) {
if !isTerminalStatus(lifecycle.Status) {
s.engine.startDeadlineCoordinator(ctx, t)
}
}
// resume the active task's engine; other tasks resume when opened (setActive).
// (a restored-paused task's loops still start but idle until resumed.)
// 排队中的任务不在此启动——留给 reconcileConcurrency 按空位补位。
if t := m.ActiveTask(); t != nil && !t.Queued {
s.engine.Run(ctx, t)
// Restore every task that had already been admitted before shutdown. Starting
// only the active UI task left other non-queued tasks counted as concurrency
// occupants without live loops, which could permanently block the persistent
// FIFO. Paused loops remain idle; queued tasks are admitted below as slots allow.
for _, t := range m.List() {
lifecycle := t.lifecycleSnapshot()
if !lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
s.engine.Run(ctx, t)
}
}
go s.reconcileConcurrency()
return s
}
@@ -636,7 +642,12 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/tasks", s.listTasks)
mux.HandleFunc("POST /api/tasks", s.createTask)
mux.HandleFunc("GET /api/task-templates", s.pgListTaskTemplates)
mux.HandleFunc("POST /api/task-templates", s.pgCreateTaskTemplate)
mux.HandleFunc("PATCH /api/task-templates/{id}", s.pgUpdateTaskTemplate)
mux.HandleFunc("DELETE /api/task-templates/{id}", s.pgDeleteTaskTemplate)
mux.HandleFunc("GET /api/tasks/{id}", s.getTask)
mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch)
mux.HandleFunc("GET /api/tasks/{id}/coverage", s.taskCoverage)
mux.HandleFunc("GET /api/tasks/{id}/coverage-graph", s.taskCoverageGraph)
mux.HandleFunc("GET /api/tasks/{id}/asset-refs", s.taskAssetRefs)
@@ -654,6 +665,7 @@ 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/{iid}/rerun", s.rerunIntent) // 重跑单条 blocked/exhausted/stopped 意图
mux.HandleFunc("POST /api/tasks/{id}/intents/rerun-blocked", s.rerunBlocked) // 批量重跑本任务全部 blocked 意图
@@ -679,10 +691,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 +729,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 +963,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 +996,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 +1016,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 +1060,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 +1092,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 +1111,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 +1131,47 @@ func (s *Server) rerunBlocked(w http.ResponseWriter, r *http.Request) {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法重跑意图")
return
}
defer s.engine.decInflight(t.ID)
intents, err := t.Store.ListByKind(db.KindIntent, 1000000)
if err != nil {
writeErr(w, 500, err.Error())
return
}
before := make([]*db.Node, 0)
for _, intent := range intents {
if intent.State == "blocked" {
copy := *intent
before = append(before, &copy)
}
}
n, err := t.Store.ReopenBlockedIntents()
if err != nil {
writeErr(w, 500, err.Error())
return
}
queued := false
if n > 0 {
s.reviveTask(t)
queued, err = s.admitTask(t, "resume")
if err != nil {
rollbackErrors := make([]string, 0)
for _, intent := range before {
if rollbackErr := restoreRerunIntent(t, intent); rollbackErr != nil {
rollbackErrors = append(rollbackErrors, fmt.Sprintf("intent %d: %v", intent.ID, rollbackErr))
}
}
if len(rollbackErrors) > 0 {
err = fmt.Errorf("%w; restore blocked intents after admission failure: %s", err, strings.Join(rollbackErrors, "; "))
}
writeErr(w, 500, err.Error())
return
}
log.Printf("[task] #%s 批量重开 %d 条 blocked 意图", t.ID, n)
}
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": n})
writeJSON(w, 200, map[string]any{"id": t.ID, "reopened": n, "queued": queued})
}
// getLLM returns the current LLM config (key never exposed).
@@ -1294,6 +1303,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 +1347,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 +1392,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 +1423,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 +1760,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 +2046,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 +2521,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 +2533,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 +3248,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.
+41 -3
View File
@@ -77,6 +77,44 @@ func canonicalTaskID(raw string) (string, bool) {
return strconv.FormatInt(n, 10), true
}
// beginTaskDelete serializes the delete barrier with pause, resume, admission,
// and FIFO reconciliation. Every lifecycle path rechecks deleting after taking
// concMu, so none can commit a contradictory state once this returns.
func (s *Server) beginTaskDelete(taskID string) bool {
s.concMu.Lock()
defer s.concMu.Unlock()
if s.engine.IsDeleting(taskID) {
return false
}
return s.engine.BeginDelete(taskID)
}
// abortTaskDelete restores the execution barrier from the committed task row,
// not from the in-memory state observed before deletion began. Keeping the task
// paused on an unreadable row is conservative: a later explicit resume can
// safely recover it without allowing work to escape an uncertain delete.
func (s *Server) abortTaskDelete(taskID string) {
s.concMu.Lock()
defer s.concMu.Unlock()
keepPaused := true
if id, err := strconv.ParseInt(taskID, 10, 64); err == nil && s.m != nil && s.m.pg != nil {
if persisted, getErr := s.m.pg.GetTask(id); getErr == nil && persisted != nil {
keepPaused = persisted.Paused || persisted.Queued
} else if task, ok := s.m.Task(taskID); ok {
state := task.lifecycleSnapshot()
keepPaused = state.Paused || state.Queued
if getErr != nil {
log.Printf("[task-delete] task %s 读取持久状态失败,使用内存状态恢复屏障: %v", taskID, getErr)
}
} else if getErr == nil {
// The request targeted a task that does not exist. Do not retain a
// synthetic pause entry after releasing its temporary delete barrier.
keepPaused = false
}
}
s.engine.AbortDelete(taskID, keepPaused)
}
func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
id, ok := canonicalTaskID(r.PathValue("id"))
if !ok {
@@ -88,14 +126,14 @@ func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "invalid JSON: "+err.Error())
return
}
if !s.engine.BeginDelete(id) {
if !s.beginTaskDelete(id) {
writeErr(w, http.StatusConflict, "任务正在删除")
return
}
deleted := false
defer func() {
if !deleted {
s.engine.AbortDelete(id)
s.abortTaskDelete(id)
}
}()
@@ -1754,7 +1792,7 @@ func (s *Server) restoreTasksAfterProfileDelete(pg *db.DB) {
// or clear the explicit chain and restore its Agent/global fallback.
// Resume only intents that were blocked by the exhausted chain; a
// paused task remains paused and terminal tasks remain immutable.
if !isTerminalStatus(task.Status) && s.taskRuntimeAvailable(task, "planner", "worker") {
if !isTerminalStatus(task.lifecycleSnapshot().Status) && s.taskRuntimeAvailable(task, "planner", "worker") {
if _, reopenErr := task.Store.ReopenIntentsByBlockedReason(db.IntentBlockedLLMQuota); reopenErr != nil {
log.Printf("[llm-profile] task %s reopen quota-blocked intents after profile delete: %v", task.ID, reopenErr)
}
+822
View File
@@ -0,0 +1,822 @@
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)
}
}
// An Agent binding outranks the task's LLM chain; roles left unbound keep using
// the chain, so binding one role does not move the others off it.
func TestTaskLLMResolutionPrefersAgentBindingOverTaskChain(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
stamp := time.Now().UnixNano()
boundID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("bound-resolution-%d", stamp), Format: "openai",
Model: "bound-model", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
chainID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("chain-resolution-%d", stamp), Format: "openai",
Model: "chain-model", APIKey: "test-key",
})
if err != nil {
_ = m.pg.DeleteProfile(boundID)
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("binding precedence", "prefer the bound profile", db.TaskCreateOptions{
LLMProfileIDs: []int64{chainID},
})
if err != nil {
_ = m.pg.DeleteProfile(boundID)
_ = m.pg.DeleteProfile(chainID)
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
_ = m.pg.SetAgentLLMProfile("worker", nil)
_ = m.pg.DeleteProfile(boundID)
_ = m.pg.DeleteProfile(chainID)
}()
s := &Server{m: m, provByProfile: map[int64]*provEntry{}}
got, err := s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if got.Source != "task_chain" || got.Model != "chain-model" {
t.Fatalf("resolution=%+v, want the task chain while no binding exists", got)
}
if err := m.pg.SetAgentLLMProfile("worker", &boundID); err != nil {
t.Fatal(err)
}
got, err = s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if !got.Available || got.Source != "agent_binding" || got.Model != "bound-model" {
t.Fatalf("resolution=%+v, want the agent binding to outrank the task chain", got)
}
got, err = s.resolveTaskRoleLLM(task, "planner")
if err != nil {
t.Fatal(err)
}
if got.Source != "task_chain" || got.Model != "chain-model" {
t.Fatalf("resolution=%+v, want an unbound role to keep using the task chain", 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())
}
}
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)
}
}
+241
View File
@@ -0,0 +1,241 @@
package server
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"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"`
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})
}
+46
View File
@@ -0,0 +1,46 @@
package server
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestWorkerControlRoutes(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
td := t.TempDir()
s := New(context.Background(), m, td, td, td)
h := s.Handler()
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatalf("sign jwt: %v", err)
}
request := func(path string) *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"action":"pause"}`))
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
return rec
}
// The single-worker route is still registered and reaches its JSON handler.
single := request("/api/tasks/missing/intents/1/control")
if single.Code != http.StatusNotFound || !strings.Contains(single.Body.String(), `"error":"task not found"`) {
t.Fatalf("single control route unavailable: status=%d body=%s", single.Code, single.Body.String())
}
// The former batch route must fall through the mux instead of reaching a task handler.
batch := request("/api/tasks/missing/intents/control/batch")
if batch.Code != http.StatusNotFound || !strings.Contains(batch.Body.String(), "404 page not found") {
t.Fatalf("batch control route still registered: status=%d body=%s", batch.Code, batch.Body.String())
}
}
+146 -2
View File
@@ -42,7 +42,7 @@ func TestEngineDeleteBarrierRejectsNewTaskOperations(t *testing.T) {
if e.BeginDelete(taskID) {
t.Fatal("duplicate delete must not acquire the barrier")
}
e.AbortDelete(taskID)
e.AbortDelete(taskID, false)
if e.IsPaused(taskID) {
t.Fatal("aborted delete should restore a previously running task")
}
@@ -58,12 +58,130 @@ func TestAbortDeletePreservesExistingPause(t *testing.T) {
if !e.BeginDelete("42") {
t.Fatal("delete barrier should be installed")
}
e.AbortDelete("42")
e.AbortDelete("42", true)
if !e.IsPaused("42") {
t.Fatal("aborted delete must preserve a pre-existing user pause")
}
}
func TestBeginTaskDeleteSerializesWithLifecycleLock(t *testing.T) {
s := &Server{engine: NewEngine(nil)}
s.concMu.Lock()
done := make(chan bool, 1)
go func() { done <- s.beginTaskDelete("42") }()
select {
case <-done:
t.Fatal("delete barrier bypassed the lifecycle lock")
case <-time.After(20 * time.Millisecond):
}
if s.engine.IsDeleting("42") {
t.Fatal("delete barrier was installed before concMu was released")
}
s.concMu.Unlock()
select {
case ok := <-done:
if !ok {
t.Fatal("delete barrier was not installed after lifecycle lock release")
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for delete barrier")
}
s.engine.AbortDelete("42", false)
}
func TestDeletingRunningTaskKeepsConcurrencySlotUntilCommit(t *testing.T) {
task := &Task{ID: "42", Status: "running"}
s := &Server{engine: NewEngine(nil)}
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install delete barrier")
}
if !s.occupiesConcurrencySlot(task) {
t.Fatal("temporary delete pause released the running task's slot before commit")
}
s.engine.AbortDelete(task.ID, false)
}
func TestTaskLifecycleRechecksDeleteBarrierAfterConcLock(t *testing.T) {
task := &Task{ID: "42", notify: make(chan struct{}, 1)}
m := &Manager{tasks: map[string]*Task{task.ID: task}}
s := &Server{m: m, engine: NewEngine(m)}
s.concMu.Lock()
started := make(chan struct{})
pauseDone := make(chan error, 1)
go func() {
close(started)
_, err := s.applyTaskControl(task, "pause")
pauseDone <- err
}()
<-started
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install delete barrier")
}
s.concMu.Unlock()
if err := <-pauseDone; err == nil || !strings.Contains(err.Error(), "正在删除") {
t.Fatalf("pause passed a delete barrier installed while waiting for concMu: %v", err)
}
s.engine.AbortDelete(task.ID, false)
s.concMu.Lock()
started = make(chan struct{})
admitDone := make(chan error, 1)
go func() {
close(started)
_, err := s.admitTask(task, "resume")
admitDone <- err
}()
<-started
if !s.engine.BeginDelete(task.ID) {
t.Fatal("install second delete barrier")
}
s.concMu.Unlock()
if err := <-admitDone; err == nil || !strings.Contains(err.Error(), "being deleted") {
t.Fatalf("admission passed a delete barrier installed while waiting for concMu: %v", err)
}
s.engine.AbortDelete(task.ID, false)
}
func TestAbortTaskDeleteUsesPersistedPauseAndQueueState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("delete abort state", "restore persisted barrier", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := &Server{m: m, engine: NewEngine(m)}
// An engine-only pause is not authoritative. The committed task is running,
// so aborting deletion must clear this stale in-memory barrier.
s.engine.Pause(task.ID, agent.AbortPausedByUser)
if !s.beginTaskDelete(task.ID) {
t.Fatal("install first delete barrier")
}
s.abortTaskDelete(task.ID)
if s.engine.IsPaused(task.ID) {
t.Fatal("stale engine pause survived although the persisted task is running")
}
// A queued row must retain its execution barrier even when BeginDelete saw the
// engine as running. This specifically guards against restoring a stale snapshot.
if err := m.EnqueueTask(task.ID, "resume"); err != nil {
t.Fatal(err)
}
if !s.beginTaskDelete(task.ID) {
t.Fatal("install second delete barrier")
}
s.abortTaskDelete(task.ID)
if !s.engine.IsPaused(task.ID) {
t.Fatal("persisted queued task lost its execution barrier after abort")
}
}
func TestStopTaskCancelsRuntimeAndClearsLifecycleState(t *testing.T) {
e := NewEngine(nil)
const taskID = "42"
@@ -182,6 +300,32 @@ func TestWaitTaskQuiescentHonorsContext(t *testing.T) {
}
}
func TestTaskChatStatusUsesAuthoritativeBusyState(t *testing.T) {
m := &Manager{tasks: map[string]*Task{"7": {ID: "7"}}}
s := &Server{
m: m,
chatBusy: map[string]bool{"7": true},
chatCancel: map[string]context.CancelCauseFunc{},
}
req := httptest.NewRequest(http.MethodGet, "/api/tasks/7/chat/status", nil)
req.SetPathValue("id", "7")
rec := httptest.NewRecorder()
s.taskChatStatus(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String())
}
var got struct {
Running bool `json:"running"`
}
if err := json.NewDecoder(rec.Body).Decode(&got); err != nil {
t.Fatal(err)
}
if !got.Running {
t.Fatal("busy main-agent turn was reported idle")
}
}
func TestTaskChatUploadRejectsDeleteBarrierBeforeCreatingDirectory(t *testing.T) {
dataDir := t.TempDir()
m := &Manager{dir: dataDir, tasks: map[string]*Task{"7": {ID: "7"}}}
+61 -33
View File
@@ -99,6 +99,12 @@ type taskLLMStreamHooks struct {
transition func(taskLLMSelection, db.TaskLLMTransition, error)
}
// current resolves the provider this role runs on, by precedence:
// Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置。
// 绑定优先于任务链:某个角色被显式指定了模型,就一直跑在那个模型上;绑定不存在或
// 构建失败时才降级到任务链,任务链为空时再降级到全局配置。
// 返回的 profile id 只有走任务链时才非零 —— streamTaskLLM 以此判断额度错误是否
// 应该推进任务的故障转移状态(绑定/全局路径不改任务链状态,沿用既有语义)。
func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) {
taskNum, err := parseTaskID(r.taskID)
if err != nil {
@@ -113,21 +119,24 @@ func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) {
}
r.s.syncTaskLLMState(pt)
t, _ := r.s.m.Task(r.taskID)
if len(pt.LLMProfileIDs) == 0 {
prov, _, ok := r.s.fallbackProviderForAgent(r.agentKey)
if !ok {
return t, 0, pt.LLMChainRevision, nil, fmt.Errorf("task %s has no available fallback LLM provider", r.taskID)
}
if prov, _, ok := r.s.agentBindingProvider(r.agentKey); ok {
return t, 0, pt.LLMChainRevision, prov, nil
}
if pt.ActiveLLMProfileID == nil {
return t, 0, pt.LLMChainRevision, nil, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")}
if len(pt.LLMProfileIDs) > 0 {
if pt.ActiveLLMProfileID == nil {
return t, 0, pt.LLMChainRevision, nil, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")}
}
prov, _, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID)
if !ok {
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, nil, fmt.Errorf("LLM profile #%d is missing or invalid", *pt.ActiveLLMProfileID)
}
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, prov, nil
}
prov, _, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID)
prov, _, ok := r.s.globalProvider()
if !ok {
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, nil, fmt.Errorf("LLM profile #%d is missing or invalid", *pt.ActiveLLMProfileID)
return t, 0, pt.LLMChainRevision, nil, fmt.Errorf("task %s has no available fallback LLM provider", r.taskID)
}
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, prov, nil
return t, 0, pt.LLMChainRevision, prov, nil
}
func parseTaskID(id string) (int64, error) {
@@ -253,7 +262,12 @@ func streamEventCommitsOutput(event llm.StreamEvent) bool {
}
}
// CompactionWindow mirrors current()'s precedence so the context window always
// matches the provider the role will actually stream on.
func (r *taskLLMRuntime) CompactionWindow() int {
if _, cfg, ok := r.s.agentBindingProvider(r.agentKey); ok {
return cfg.CompactionWindow()
}
taskNum, err := parseTaskID(r.taskID)
if err != nil {
return (agent.Config{}).CompactionWindow()
@@ -263,7 +277,7 @@ func (r *taskLLMRuntime) CompactionWindow() int {
return (agent.Config{}).CompactionWindow()
}
if len(chain) == 0 {
if _, cfg, ok := r.s.fallbackProviderForAgent(r.agentKey); ok {
if _, cfg, ok := r.s.globalProvider(); ok {
return cfg.CompactionWindow()
}
return (agent.Config{}).CompactionWindow()
@@ -283,15 +297,27 @@ func (r *taskLLMRuntime) CompactionWindow() int {
return minimum
}
// fallbackProviderForAgent resolves the pre-existing precedence used by tasks
// without an explicit chain: the Agent's saved binding first, then the current
// global provider (which may come from the default profile or environment).
func (s *Server) fallbackProviderForAgent(agentKey string) (llm.Provider, agent.Config, bool) {
if id := s.effectiveProfileForAgent(agentKey, nil); id != nil {
if prov, cfg, ok := s.providerForProfile(*id); ok {
return s.poolForBinding(*id, prov, cfg), cfg, true
}
// agentBindingProvider resolves the profile a role is explicitly bound to
// (agents.llm_profile_id) — the highest-precedence level for task agents. ok=false
// when the role has no binding or the bound profile no longer builds, so callers
// fall through to the task chain. A bound profile stays exclusive unless
// llm_pool_bind_fallback is on, which is what poolForBinding encodes.
func (s *Server) agentBindingProvider(agentKey string) (llm.Provider, agent.Config, bool) {
id := s.effectiveProfileForAgent(agentKey, nil)
if id == nil {
return nil, agent.Config{}, false
}
prov, cfg, ok := s.providerForProfile(*id)
if !ok {
return nil, agent.Config{}, false
}
return s.poolForBinding(*id, prov, cfg), cfg, true
}
// globalProvider returns the process-wide provider (persisted active profile or
// environment config) — the last resort once a role has neither a binding nor a
// task chain.
func (s *Server) globalProvider() (llm.Provider, agent.Config, bool) {
s.cfgMu.Lock()
defer s.cfgMu.Unlock()
if !s.llmOn || s.llmProv == nil {
@@ -300,30 +326,32 @@ func (s *Server) fallbackProviderForAgent(agentKey string) (llm.Provider, agent.
return s.llmProv, s.llmCfg, true
}
// taskRuntimeAvailable reports whether every listed role can resolve a provider
// under the runtime precedence in current(): the role's own binding first, then
// the task chain, then global. An exhausted chain is a hard stop for unbound
// roles rather than a silent fall through to global — same as current().
func (s *Server) taskRuntimeAvailable(t *Task, agentKeys ...string) bool {
if t == nil {
if t == nil || len(agentKeys) == 0 {
return false
}
state := t.llmStateSnapshot()
unboundReady := false
if len(state.ProfileIDs) > 0 {
if state.ActiveID == nil {
return false
if state.ActiveID != nil {
_, _, unboundReady = s.providerForProfile(*state.ActiveID)
}
_, _, ok := s.providerForProfile(*state.ActiveID)
return ok
}
s.cfgMu.Lock()
globallyReady := s.llmOn && s.llmProv != nil
s.cfgMu.Unlock()
if globallyReady {
return true
} else {
_, _, unboundReady = s.globalProvider()
}
for _, key := range agentKeys {
if _, _, ok := s.fallbackProviderForAgent(key); !ok {
if _, _, ok := s.agentBindingProvider(key); ok {
continue
}
if !unboundReady {
return false
}
}
return len(agentKeys) > 0
return true
}
func (s *Server) invalidateTaskAgents() {
@@ -375,7 +403,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()
}
+47 -2
View File
@@ -122,7 +122,7 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
defer wg.Done()
for i := 0; i < 500; i++ {
id := int64(offset*1000 + i + 1)
task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "")
task.setLLMState(&id, &id, []int64{id, id + 1}, int64(i), "ready", "")
}
}(writer)
}
@@ -146,6 +146,24 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
}
}
func TestTaskLLMStateRejectsOlderRevision(t *testing.T) {
t.Parallel()
task := &Task{ID: "7"}
current := int64(22)
stale := int64(11)
if !task.setLLMState(&current, &current, []int64{current}, 2, "ready", "") {
t.Fatal("initial LLM state was rejected")
}
if task.setLLMState(&stale, &stale, []int64{stale}, 1, "ready", "stale") {
t.Fatal("older LLM snapshot was accepted")
}
state := task.llmStateSnapshot()
if state.ChainRevision != 2 || state.ActiveID == nil || *state.ActiveID != current ||
len(state.ProfileIDs) != 1 || state.ProfileIDs[0] != current || state.FailoverReason != "" {
t.Fatalf("older LLM snapshot overwrote current state: %+v", state)
}
}
func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) {
t.Parallel()
quota := errors.New("openai: status 402: insufficient_quota")
@@ -381,7 +399,7 @@ func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) {
engine.started.Store("7", true)
task := &Task{ID: "7", Status: "running"}
profileID := int64(11)
task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted")
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
s := &Server{engine: engine}
if got := s.resolvedTaskStatus(task); got != "running" {
t.Fatalf("exhausted started status=%q, want running", got)
@@ -392,6 +410,33 @@ func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) {
}
}
// A role with no Agent binding falls through to the task chain, then to global —
// and an exhausted chain stays a hard stop instead of silently borrowing global.
func TestTaskRuntimeAvailableUnboundRoleFollowsChainThenGlobal(t *testing.T) {
t.Parallel()
// A nil pg means effectiveProfileForAgent finds no binding for any role, so
// every role here resolves through the chain/global levels.
s := &Server{m: &Manager{}}
task := &Task{ID: "7"}
if s.taskRuntimeAvailable(task, "worker") {
t.Fatal("no binding, no chain and no global provider must not be runnable")
}
s.llmOn, s.llmProv = true, &scriptedLLMProvider{}
if !s.taskRuntimeAvailable(task, "worker") {
t.Fatal("global provider must serve a role with no binding and no chain")
}
if s.taskRuntimeAvailable(task) {
t.Fatal("a task with no roles requested must not report as runnable")
}
profileID := int64(11)
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
if s.taskRuntimeAvailable(task, "worker") {
t.Fatal("exhausted chain must not fall back to the global provider")
}
}
func TestAuthoritativeTaskResolverDoesNotFallBackWhenChainUnavailable(t *testing.T) {
t.Parallel()
e := &Engine{}
+123
View File
@@ -0,0 +1,123 @@
package server
import (
"fmt"
"net/http"
"github.com/Autumn-27/artex/db"
)
type taskLLMResolution struct {
ProfileID *int64 `json:"profile_id,omitempty"`
Name string `json:"name"`
Format string `json:"format"`
Model string `json:"model"`
Source string `json:"source"`
Available bool `json:"available"`
Reason string `json:"reason,omitempty"`
}
func (s *Server) resolutionFromProfile(p *db.LLMProfile, source string) taskLLMResolution {
if p == nil {
return taskLLMResolution{Source: source, Reason: "LLM 配置不存在"}
}
id := p.ID
result := taskLLMResolution{
ProfileID: &id,
Name: p.Name,
Format: p.Format,
Model: p.Model,
Source: source,
}
if p.APIKey == "" {
result.Reason = "LLM 配置未设置 API Key"
return result
}
if _, _, ok := s.providerForProfile(p.ID); !ok {
result.Reason = "LLM 配置格式或参数无效"
return result
}
result.Available = true
return result
}
// resolveTaskRoleLLM mirrors taskLLMRuntime.current without creating a provider:
// role Agent binding -> explicit task chain -> 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")
}
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
}
}
}
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
}
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)
}
+172
View File
@@ -0,0 +1,172 @@
package server
import (
"errors"
"fmt"
"net/http"
"strings"
"unicode/utf8"
"github.com/Autumn-27/artex/db"
)
const maxTaskTemplateRequestBytes = 512 << 10
// taskTemplateRequest uses pointers so PATCH can distinguish omitted fields from
// explicit empty values. Empty values are still rejected by the DB validator.
type taskTemplateRequest struct {
Name *string `json:"name"`
Description *string `json:"description"`
Goal *string `json:"goal"`
}
func decodeTaskTemplateRequest(w http.ResponseWriter, r *http.Request, req *taskTemplateRequest) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxTaskTemplateRequestBytes)
if err := decode(r, req); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, err.Error())
}
return false
}
return true
}
func validateTaskTemplateRequest(req taskTemplateRequest) error {
checks := []struct {
name string
value *string
limit int
}{
{name: "name", value: req.Name, limit: db.MaxTaskTemplateNameRunes},
{name: "description", value: req.Description, limit: db.MaxTaskTemplateTextRunes},
{name: "goal", value: req.Goal, limit: db.MaxTaskTemplateTextRunes},
}
for _, check := range checks {
if check.value == nil {
continue
}
value := strings.TrimSpace(*check.value)
if check.name == "name" {
value = strings.Join(strings.Fields(value), " ")
}
if utf8.RuneCountInString(value) > check.limit {
return fmt.Errorf("%s 最多 %d 个字符", check.name, check.limit)
}
}
return nil
}
func writeTaskTemplateErr(w http.ResponseWriter, err error) {
switch {
case errors.Is(err, db.ErrTaskTemplateInvalid):
writeErr(w, http.StatusBadRequest, err.Error())
case errors.Is(err, db.ErrTaskTemplateNameConflict):
writeErr(w, http.StatusConflict, "模板名称已存在")
case errors.Is(err, db.ErrTaskTemplateNotFound):
writeErr(w, http.StatusNotFound, "task template not found")
default:
writeErr(w, http.StatusInternalServerError, err.Error())
}
}
func (s *Server) pgListTaskTemplates(w http.ResponseWriter, _ *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
templates, err := pg.ListTaskTemplates()
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]any{"templates": templates})
}
func (s *Server) pgCreateTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
var req taskTemplateRequest
if !decodeTaskTemplateRequest(w, r, &req) {
return
}
if err := validateTaskTemplateRequest(req); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
template, err := pg.CreateTaskTemplate(db.TaskTemplateInput{
Name: stringValue(req.Name),
Description: stringValue(req.Description),
Goal: stringValue(req.Goal),
})
if err != nil {
writeTaskTemplateErr(w, err)
return
}
writeJSON(w, http.StatusCreated, template)
}
func (s *Server) pgUpdateTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, http.StatusBadRequest, "bad task template id")
return
}
var req taskTemplateRequest
if !decodeTaskTemplateRequest(w, r, &req) {
return
}
if err := validateTaskTemplateRequest(req); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
if req.Name == nil && req.Description == nil && req.Goal == nil {
writeErr(w, http.StatusBadRequest, "至少需要提供 name、description 或 goal")
return
}
template, err := pg.PatchTaskTemplate(id, db.TaskTemplatePatch{
Name: req.Name, Description: req.Description, Goal: req.Goal,
})
if err != nil {
writeTaskTemplateErr(w, err)
return
}
writeJSON(w, http.StatusOK, template)
}
func (s *Server) pgDeleteTaskTemplate(w http.ResponseWriter, r *http.Request) {
pg := s.pg(w)
if pg == nil {
return
}
id, ok := pathInt(r, "id")
if !ok {
writeErr(w, http.StatusBadRequest, "bad task template id")
return
}
deleted, err := pg.DeleteTaskTemplate(id)
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
if !deleted {
writeErr(w, http.StatusNotFound, "task template not found")
return
}
writeJSON(w, http.StatusOK, map[string]any{"deleted": id})
}
func stringValue(value *string) string {
if value == nil {
return ""
}
return *value
}
+163
View File
@@ -0,0 +1,163 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestTaskTemplateHTTPCRUD(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
h := s.Handler()
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
do := func(method, path string, body any) (int, map[string]any) {
var data []byte
if body != nil {
data, _ = json.Marshal(body)
}
req := httptest.NewRequest(method, path, bytes.NewReader(data))
req.Header.Set("Authorization", "Bearer "+token)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
var out map[string]any
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out
}
name := fmt.Sprintf("http-template-%d", time.Now().UnixNano())
code, created := do(http.MethodPost, "/api/task-templates", map[string]any{
"name": name, "description": "description", "goal": "goal",
})
if code != http.StatusCreated {
t.Fatalf("create status=%d body=%v", code, created)
}
id := int64(created["id"].(float64))
t.Cleanup(func() { _, _ = m.pg.DeleteTaskTemplate(id) })
code, duplicate := do(http.MethodPost, "/api/task-templates", map[string]any{
"name": " " + name + " ", "description": "other", "goal": "other",
})
if code != http.StatusConflict {
t.Fatalf("duplicate status=%d body=%v", code, duplicate)
}
code, updated := do(http.MethodPatch, fmt.Sprintf("/api/task-templates/%d", id), map[string]any{"goal": "patched goal"})
if code != http.StatusOK || updated["goal"] != "patched goal" || updated["description"] != "description" {
t.Fatalf("patch status=%d body=%v", code, updated)
}
code, _ = do(http.MethodPatch, fmt.Sprintf("/api/task-templates/%d", id), map[string]any{
"goal": strings.Repeat("界", db.MaxTaskTemplateTextRunes+1),
})
if code != http.StatusBadRequest {
t.Fatalf("overlong patch status=%d, want %d", code, http.StatusBadRequest)
}
code, _ = do(http.MethodPost, "/api/task-templates", map[string]any{
"name": "oversized", "description": strings.Repeat("x", maxTaskTemplateRequestBytes+1), "goal": "goal",
})
if code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized create status=%d, want %d", code, http.StatusRequestEntityTooLarge)
}
code, listed := do(http.MethodGet, "/api/task-templates", nil)
if code != http.StatusOK || listed["templates"] == nil {
t.Fatalf("list status=%d body=%v", code, listed)
}
code, deleted := do(http.MethodDelete, fmt.Sprintf("/api/task-templates/%d", id), nil)
if code != http.StatusOK || int64(deleted["deleted"].(float64)) != id {
t.Fatalf("delete status=%d body=%v", code, deleted)
}
}
func TestConversationPatchReturnsPinState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer m.Close()
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
conversation, err := m.pg.CreateConversation("mainagent", "pin through http", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = m.pg.DeleteConversation(conversation.ID) })
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
body, _ := json.Marshal(map[string]any{"title": " pinned title ", "pinned": true})
req := httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec := httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("patch status=%d body=%s", rec.Code, rec.Body.String())
}
var updated struct {
Title string `json:"title"`
Pinned bool `json:"pinned"`
PinnedAt *string `json:"pinned_at"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &updated); err != nil {
t.Fatal(err)
}
if updated.Title != "pinned title" || !updated.Pinned || updated.PinnedAt == nil {
t.Fatalf("unexpected conversation patch response: %+v", updated)
}
body, _ = json.Marshal(map[string]any{"title": strings.Repeat("会", maxConversationTitleRunes+1)})
req = httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("overlong conversation title status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{"title": strings.Repeat("x", maxConversationRequestBytes+1)})
req = httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/api/conversations/%d", conversation.ID), bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized conversation patch status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{
"agent_key": "mainagent", "title": strings.Repeat("会", maxConversationTitleRunes+1),
})
req = httptest.NewRequest(http.MethodPost, "/api/conversations", bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("overlong conversation create status=%d body=%s", rec.Code, rec.Body.String())
}
body, _ = json.Marshal(map[string]any{
"agent_key": "mainagent", "title": strings.Repeat("x", maxConversationRequestBytes+1),
})
req = httptest.NewRequest(http.MethodPost, "/api/conversations", bytes.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
rec = httptest.NewRecorder()
s.Handler().ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("oversized conversation create status=%d body=%s", rec.Code, rec.Body.String())
}
}
+21 -1
View File
@@ -219,7 +219,7 @@ func (t *Traffic) record(f *mproxy.Flow) {
respBody := t.bodyOrBlob(f.Response.Body)
ct := f.Response.Header.Get("Content-Type")
reqTxt := fmt.Sprintf("%s %s %s\n%s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, headerLines(f.Request.Header), reqBody)
reqTxt := fmt.Sprintf("%s %s %s\n%s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, requestHeaderLines(f.Request), reqBody)
respTxt := fmt.Sprintf("HTTP %d\n%s\n%s", f.Response.StatusCode, headerLines(f.Response.Header), respBody)
_ = os.WriteFile(filepath.Join(exDir, "request.http"), []byte(reqTxt), 0o644)
_ = os.WriteFile(filepath.Join(exDir, "response.http"), []byte(respTxt), 0o644)
@@ -267,6 +267,26 @@ func headerLines(h map[string][]string) string {
return b.String()
}
// requestHeaderLines restores the HTTP Host header, which net/http stores on
// Request.Host rather than in Header. Keeping it in request.http makes the raw
// capture complete and directly replayable.
func requestHeaderLines(req *mproxy.Request) string {
headers := req.Header.Clone()
headers.Del("Host")
host := req.URL.Host
if raw := req.Raw(); raw != nil && strings.TrimSpace(raw.Host) != "" {
host = raw.Host
}
var b strings.Builder
if host = strings.TrimSpace(host); host != "" {
b.WriteString("Host: ")
b.WriteString(host)
b.WriteByte('\n')
}
b.WriteString(headerLines(headers))
return b.String()
}
func sanitize(s string) string {
r := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "?", "_", "*", "_", "\"", "_", "<", "_", ">", "_", "|", "_")
out := r.Replace(s)
+18
View File
@@ -2,12 +2,30 @@ package traffic
import (
"fmt"
"net/http"
"net/url"
"os"
"path/filepath"
"strings"
"testing"
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
)
func TestRequestHeaderLinesIncludesHost(t *testing.T) {
req := &mproxy.Request{
URL: &url.URL{Host: "target.example:8443"},
Header: http.Header{"Accept": []string{"application/json"}},
}
got := requestHeaderLines(req)
if !strings.Contains(got, "Host: target.example:8443\n") {
t.Fatalf("request headers missing Host: %q", got)
}
if !strings.Contains(got, "Accept: application/json\n") {
t.Fatalf("request headers missing regular header: %q", got)
}
}
// TestDeleteHost verifies the delete contract: rows for hosts containing the
// substring are removed together with their file trees, non-matching hosts are
// untouched, and the count is right.
@@ -35,7 +35,7 @@ export function MainContent({ children }: { children: ReactNode }) {
useEffect(() => {
api
.health()
.then((h) => setVersion(h.version ?? ""))
.then((h) => setVersion((h.version ?? "").replace(/^v(?=\d)/, "")))
.catch(() => setVersion(""));
}, []);
+195 -80
View File
@@ -1,20 +1,27 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { Bot, ChevronDownIcon, Loader2Icon, PaperclipIcon, PlusIcon, SendIcon, Square, Trash2Icon, XIcon, ZapIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Textarea } from "@/components/ui/textarea";
import { Badge } from "@/components/ui/badge";
import { ScrollArea } from "@/components/ui/scroll-area";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
ArrowUpIcon,
Bot,
ChevronDownIcon,
Loader2Icon,
MoreHorizontalIcon,
PaperclipIcon,
PencilIcon,
PinIcon,
PinOffIcon,
PlusIcon,
Square,
Trash2Icon,
XIcon,
ZapIcon,
} from "lucide-react";
import { toast } from "sonner";
import { TodoPopover } from "@/components/todo-popover";
import { Transcript } from "@/components/transcript";
import {
AlertDialog,
AlertDialogAction,
@@ -24,15 +31,21 @@ import {
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
AlertDialogTrigger,
} from "@/components/ui/alert-dialog";
import { Transcript } from "@/components/transcript";
import { TodoPopover } from "@/components/todo-popover";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
DropdownMenu,
DropdownMenuContent,
DropdownMenuGroup,
DropdownMenuItem,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import { api } from "@/lib/api";
import type { Activity, Agent, ChatAttachment, Conversation, LLMProfile } from "@/lib/types";
import { cn } from "@/lib/utils";
@@ -46,8 +59,8 @@ function fmtBytes(n: number): string {
// fmtTokens renders a compact token count (1234 → 1.2k, 2_000_000 → 2M).
function fmtTokens(n: number): string {
if (n >= 1_000_000) return (n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1) + "M";
if (n >= 1000) return (n / 1000).toFixed(n >= 10000 ? 0 : 1) + "k";
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(n >= 10_000_000 ? 0 : 1)}M`;
if (n >= 1000) return `${(n / 1000).toFixed(n >= 10000 ? 0 : 1)}k`;
return String(n);
}
@@ -66,6 +79,10 @@ function fmtDuration(ms: number): string {
// thread stays snappy (only ~a page of rows is in the DOM until you scroll up).
const HISTORY_PAGE = 200;
function conversationIsPinned(conversation: Conversation): boolean {
return conversation.pinned ?? Boolean(conversation.pinned_at);
}
// LiveBadge is the small pulsing "实时" chip reused from the task's main-agent
// console — shown while a turn is streaming.
function LiveBadge() {
@@ -113,7 +130,7 @@ function Composer({
const atts = attachments ?? [];
function onKeyDown(e: React.KeyboardEvent<HTMLTextAreaElement>) {
// Enter 发送;Shift+Enter 换行。isComposing 时是中文等输入法在选字,不拦截。
if (e.key !== "Enter" || e.shiftKey || e.nativeEvent.isComposing) return;
if (e.key !== "Enter" || e.shiftKey || e.nativeEvent.isComposing || e.nativeEvent.keyCode === 229) return;
e.preventDefault();
onSend();
}
@@ -144,8 +161,8 @@ function Composer({
))}
</div>
)}
<div className="flex items-end gap-2">
{leftSlot}
<div className="flex flex-wrap items-end gap-2">
{leftSlot ? <div className="w-full sm:w-auto">{leftSlot}</div> : null}
{onPickFiles && (
<>
<input
@@ -170,7 +187,7 @@ function Composer({
</>
)}
<Textarea
className="max-h-40 min-h-10 flex-1 resize-none"
className="max-h-40 min-h-10 min-w-0 flex-1 resize-none"
rows={1}
placeholder={placeholder}
value={value}
@@ -185,8 +202,14 @@ function Composer({
<Square className="size-3.5 fill-current" />
</Button>
) : (
<Button size="icon" onClick={onSend} disabled={disabled || (!value.trim() && atts.length === 0)}>
<SendIcon className="size-4" />
<Button
size="icon"
onClick={onSend}
disabled={disabled || (!value.trim() && atts.length === 0)}
title="发送消息"
aria-label="发送消息"
>
<ArrowUpIcon />
</Button>
)}
</div>
@@ -234,7 +257,10 @@ function LLMProfileRow({
{/* default option */}
<button
type="button"
onClick={() => { onChange(null); setOpen(false); }}
onClick={() => {
onChange(null);
setOpen(false);
}}
className={cn(
"flex w-full flex-col rounded px-2 py-1.5 text-left hover:bg-accent",
selected == null && "bg-accent",
@@ -242,21 +268,28 @@ function LLMProfileRow({
>
<span className="text-sm">默认{activeDefault ? `(${activeDefault.name})` : ""}</span>
{activeDefault && (
<span className="text-muted-foreground text-[11px]">{activeDefault.format} · {activeDefault.model}</span>
<span className="text-muted-foreground text-[11px]">
{activeDefault.format} · {activeDefault.model}
</span>
)}
</button>
{profiles.map((p) => (
<button
key={p.id}
type="button"
onClick={() => { onChange(Number(p.id)); setOpen(false); }}
onClick={() => {
onChange(Number(p.id));
setOpen(false);
}}
className={cn(
"flex w-full flex-col rounded px-2 py-1.5 text-left hover:bg-accent",
selected === Number(p.id) && "bg-accent",
)}
>
<span className="text-sm">{p.name}</span>
<span className="text-muted-foreground text-[11px]">{p.format} · {p.model}</span>
<span className="text-muted-foreground text-[11px]">
{p.format} · {p.model}
</span>
</button>
))}
</PopoverContent>
@@ -326,7 +359,7 @@ function DraftChat({
const agentPicker = (
<Select value={agentKey} onValueChange={setAgentKey}>
<SelectTrigger size="sm" className="w-40 shrink-0">
<SelectTrigger size="sm" className="w-full sm:w-40">
<SelectValue placeholder="选择 Agent…" />
</SelectTrigger>
<SelectContent>
@@ -336,7 +369,9 @@ function DraftChat({
<Bot className="size-3.5" />
{a.name}
{!a.builtin && (
<Badge variant="outline" className="px-1 py-0 text-[9px]">自定义</Badge>
<Badge variant="outline" className="px-1 py-0 text-[9px]">
自定义
</Badge>
)}
</span>
</SelectItem>
@@ -353,9 +388,7 @@ function DraftChat({
<Bot className="text-primary size-6" />
</div>
<div className="text-sm font-medium">开始和「{agent?.name ?? "Agent"}」对话</div>
{agent?.description && (
<p className="text-muted-foreground max-w-md text-xs">{agent.description}</p>
)}
{agent?.description && <p className="text-muted-foreground max-w-md text-xs">{agent.description}</p>}
</div>
<Composer
@@ -424,10 +457,7 @@ function ChatView({
}
// conversation-scoped detail fetcher for the reused Transcript renderer.
const fetchDetail = React.useCallback(
(seq: number) => api.conversationMsgDetail(conv.id, seq),
[conv.id],
);
const fetchDetail = React.useCallback((seq: number) => api.conversationMsgDetail(conv.id, seq), [conv.id]);
// seq of the most-recent TodoWrite tool call (for the Todo popover); null if none.
const latestTodoSeq = React.useMemo(() => {
@@ -462,7 +492,9 @@ function ChatView({
setHasMore(r.hasMore);
setRunning(r.running);
})
.catch(() => {});
.catch(() => {
/* The empty state remains usable when history loading fails. */
});
return () => {
live = false;
};
@@ -497,9 +529,7 @@ function ChatView({
const contentRef = React.useRef<HTMLDivElement | null>(null);
const atBottomRef = React.useRef(true);
const viewport = React.useCallback(
() =>
(contentRef.current?.closest('[data-slot="scroll-area-viewport"]') as HTMLElement | null) ??
null,
() => (contentRef.current?.closest('[data-slot="scroll-area-viewport"]') as HTMLElement | null) ?? null,
[],
);
// Scroll-up loads one older page and prepends it, preserving the visual position
@@ -543,6 +573,7 @@ function ChatView({
return () => vp.removeEventListener("scroll", onScroll);
}, [viewport, loadEarlier]);
// open/switch a conversation → jump to the latest (bottom)
// biome-ignore lint/correctness/useExhaustiveDependencies: changing conversations intentionally retriggers the scroll reset.
React.useLayoutEffect(() => {
const vp = viewport();
if (vp) {
@@ -551,6 +582,7 @@ function ChatView({
}
}, [conv.id, viewport]);
// new activity → stick to bottom only if the user is already pinned there
// biome-ignore lint/correctness/useExhaustiveDependencies: message and running changes intentionally retrigger bottom anchoring.
React.useLayoutEffect(() => {
if (!atBottomRef.current) return;
const vp = viewport();
@@ -560,8 +592,12 @@ function ChatView({
// Per-conversation token total, live — same accounting as the main-agent
// console: completed runs' `result` sum + the in-progress run's latest `usage`.
const tokenTotal = React.useMemo(() => {
let i = 0, o = 0, cr = 0;
let li = 0, lo = 0, lcr = 0;
let i = 0,
o = 0,
cr = 0;
let li = 0,
lo = 0,
lcr = 0;
let turns = 0; // agent 循环轮次 = 模型调用次数(每次一条 kind='usage')
for (const a of messages) {
if (a.kind === "result") {
@@ -576,7 +612,9 @@ function ChatView({
lcr = a.cache_read_tokens ?? 0;
}
}
const I = i + li, O = o + lo, CR = cr + lcr;
const I = i + li,
O = o + lo,
CR = cr + lcr;
return { i: I, o: O, cr: CR, turns, any: I + O + CR > 0 };
}, [messages]);
@@ -636,10 +674,10 @@ function ChatView({
return (
<>
{/* header: which agent + live + token meta */}
<div className="flex items-center gap-2 border-b px-4 py-2.5">
<div className="flex min-w-0 flex-wrap items-center gap-2 border-b px-4 py-2.5">
<Bot className="text-muted-foreground size-4 shrink-0" />
<span className="shrink-0 text-sm font-medium">{agent?.name ?? conv.agent_key}</span>
<span className="text-muted-foreground shrink-0 font-mono text-xs">{conv.agent_key}</span>
<span className="min-w-0 max-w-48 truncate text-sm font-medium">{agent?.name ?? conv.agent_key}</span>
<span className="text-muted-foreground hidden shrink-0 font-mono text-xs sm:inline">{conv.agent_key}</span>
{agent && !agent.builtin && (
<Badge variant="outline" className="shrink-0 px-1.5 py-0 text-[10px]">
自定义
@@ -649,14 +687,14 @@ function ChatView({
<span className="text-muted-foreground min-w-0 truncate text-xs">{agent.description}</span>
)}
{running && <LiveBadge />}
<div className="text-muted-foreground ml-auto flex shrink-0 items-center gap-3 text-xs">
<div className="text-muted-foreground ml-auto flex min-w-0 max-w-full items-center justify-end gap-x-3 gap-y-1 text-xs max-sm:w-full max-sm:flex-wrap">
{tokenTotal.turns > 0 && (
<span title="agent 循环轮次(模型调用次数)" className="tabular-nums">
{tokenTotal.turns} 轮
</span>
)}
{tokenTotal.any && (
<span title="input / cache(read) / output tokens" className="tabular-nums">
<span title="input / cache(read) / output tokens" className="min-w-0 truncate tabular-nums">
input {fmtTokens(tokenTotal.i)} · cache {fmtTokens(tokenTotal.cr)} · output {fmtTokens(tokenTotal.o)}
</span>
)}
@@ -664,10 +702,7 @@ function ChatView({
</div>
{/* messages */}
<ScrollArea
type="auto"
className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!"
>
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<div className="min-w-0 max-w-full px-4 py-3" ref={contentRef}>
{messages.length === 0 && !running ? (
<div className="text-muted-foreground py-10 text-center text-sm">
@@ -676,9 +711,7 @@ function ChatView({
) : (
<>
{hasMore && (
<div className="text-muted-foreground/70 pb-2 text-center text-[11px]">
向上滚动加载更早的消息…
</div>
<div className="text-muted-foreground/70 pb-2 text-center text-[11px]">向上滚动加载更早的消息…</div>
)}
<Transcript activity={messages} live={running} chat fetchDetail={fetchDetail} />
</>
@@ -711,8 +744,8 @@ function ChatView({
);
}
// ConversationItem is one row in the left rail — status-ish icon, title + agent
// subtitle, inline rename, hover delete. Modeled on the console's SessionItem.
// ConversationItem is one row in the left rail: title, agent subtitle, inline
// rename, pin marker, and a compact action menu.
function ConversationItem({
conv,
agent,
@@ -723,6 +756,8 @@ function ConversationItem({
onStartRename,
onRenameText,
onCommitRename,
onCancelRename,
onTogglePinned,
onDelete,
}: {
conv: Conversation;
@@ -734,8 +769,22 @@ function ConversationItem({
onStartRename: () => void;
onRenameText: (v: string) => void;
onCommitRename: () => void;
onCancelRename: () => void;
onTogglePinned: () => void;
onDelete: () => void;
}) {
const [deleteOpen, setDeleteOpen] = React.useState(false);
const renameInputRef = React.useRef<HTMLInputElement>(null);
const cancelRenameRef = React.useRef(false);
const pinned = conversationIsPinned(conv);
React.useEffect(() => {
if (!renaming) return;
cancelRenameRef.current = false;
renameInputRef.current?.focus();
renameInputRef.current?.select();
}, [renaming]);
return (
<div
className={cn(
@@ -745,13 +794,19 @@ function ConversationItem({
>
{renaming ? (
<input
autoFocus
ref={renameInputRef}
value={renameText}
onChange={(e) => onRenameText(e.target.value)}
onBlur={onCommitRename}
onBlur={() => {
if (!cancelRenameRef.current) onCommitRename();
}}
onKeyDown={(e) => {
if (e.key === "Enter") onCommitRename();
if (e.key === "Escape") onCommitRename();
if (e.key === "Escape") {
e.preventDefault();
cancelRenameRef.current = true;
onCancelRename();
}
}}
className="border-input bg-background min-w-0 flex-1 rounded-md border px-2 py-1 text-sm"
/>
@@ -763,26 +818,58 @@ function ConversationItem({
title="双击重命名"
className="min-w-0 flex-1 rounded-md px-2 py-1.5 text-left"
>
<div className="truncate text-sm">{conv.title || "新对话"}</div>
<div className="flex min-w-0 items-center gap-1.5">
{pinned && <PinIcon className="text-primary size-3 shrink-0" aria-label="已置顶" />}
<div className="truncate text-sm">{conv.title || "新对话"}</div>
</div>
<div className="text-muted-foreground flex min-w-0 items-center gap-1 text-[11px]">
<Bot className="size-3 shrink-0" />
<span className="min-w-0 truncate">{agent?.name ?? conv.agent_key}</span>
<span className="shrink-0">·</span>
<span className="shrink-0">{new Date(conv.created_at).toLocaleDateString("zh-CN", { month: "numeric", day: "numeric", hour: "2-digit", minute: "2-digit" })}</span>
<span className="shrink-0">
{new Date(conv.created_at).toLocaleDateString("zh-CN", {
month: "numeric",
day: "numeric",
hour: "2-digit",
minute: "2-digit",
})}
</span>
<span className="shrink-0 opacity-60">#{conv.id}</span>
</div>
</button>
)}
<AlertDialog>
<AlertDialogTrigger asChild>
<DropdownMenu>
<DropdownMenuTrigger asChild>
<Button
variant="ghost"
size="icon-sm"
className="text-muted-foreground hover:text-destructive shrink-0 opacity-0 transition-opacity group-hover:opacity-100"
className="text-muted-foreground shrink-0"
aria-label={`管理对话「${conv.title || "新对话"}」`}
>
<Trash2Icon className="size-3.5" />
<MoreHorizontalIcon />
</Button>
</AlertDialogTrigger>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuGroup>
<DropdownMenuItem onSelect={onStartRename}>
<PencilIcon />
重命名
</DropdownMenuItem>
<DropdownMenuItem onSelect={onTogglePinned}>
{pinned ? <PinOffIcon /> : <PinIcon />}
{pinned ? "取消置顶" : "置顶"}
</DropdownMenuItem>
</DropdownMenuGroup>
<DropdownMenuSeparator />
<DropdownMenuGroup>
<DropdownMenuItem variant="destructive" onSelect={() => setDeleteOpen(true)}>
<Trash2Icon />
删除
</DropdownMenuItem>
</DropdownMenuGroup>
</DropdownMenuContent>
</DropdownMenu>
<AlertDialog open={deleteOpen} onOpenChange={setDeleteOpen}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>删除对话「{conv.title || "新对话"}」?</AlertDialogTitle>
@@ -813,11 +900,24 @@ export default function ChatPage() {
>({});
const reloadConvs = React.useCallback(() => {
api.conversations().then(setConvs).catch(() => setConvs([]));
api
.conversations()
.then(setConvs)
.catch(() => setConvs([]));
}, []);
React.useEffect(() => {
api.agents().then(setAgents).catch(() => {});
api.llmProfiles().then(setProfiles).catch(() => {});
api
.agents()
.then(setAgents)
.catch(() => {
/* Agent metadata is optional for rendering an existing conversation. */
});
api
.llmProfiles()
.then(setProfiles)
.catch(() => {
/* The conversation remains usable without profile labels. */
});
reloadConvs();
}, [reloadConvs]);
@@ -856,6 +956,16 @@ export default function ChatPage() {
}
}
async function togglePinned(conversation: Conversation) {
const pinned = conversationIsPinned(conversation);
try {
await api.pinConversation(conversation.id, !pinned);
reloadConvs();
} catch (e) {
toast.error(`${pinned ? "取消置顶" : "置顶"}失败:${(e as Error).message}`);
}
}
function startRename(c: Conversation) {
setRenamingId(c.id);
setRenameText(c.title || "");
@@ -876,9 +986,9 @@ export default function ChatPage() {
return (
<div
data-content-padding="false"
className="flex h-[calc(100svh-3rem)] flex-col overflow-hidden p-4 md:h-[calc(100svh-4rem)] md:p-6"
className="flex h-[calc(100svh-3rem)] min-w-0 flex-col overflow-hidden p-3 sm:p-4 md:h-[calc(100svh-4rem)] md:p-6"
>
<div className="grid min-h-0 flex-1 grid-cols-[18rem_1fr] grid-rows-[minmax(0,1fr)] gap-4">
<div className="grid min-h-0 min-w-0 flex-1 grid-cols-1 grid-rows-[minmax(10rem,15rem)_minmax(0,1fr)] gap-3 md:grid-cols-[18rem_minmax(0,1fr)] md:grid-rows-[minmax(0,1fr)] md:gap-4">
{/* left: conversation list */}
<div className="bg-card flex flex-col overflow-hidden rounded-lg border">
<div className="border-b p-2">
@@ -888,9 +998,7 @@ export default function ChatPage() {
</div>
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<div className="flex min-w-0 flex-col gap-0.5 p-2">
{convs.length === 0 && (
<p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>
)}
{convs.length === 0 && <p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>}
{convs.map((c) => (
<ConversationItem
key={c.id}
@@ -903,6 +1011,8 @@ export default function ChatPage() {
onStartRename={() => startRename(c)}
onRenameText={setRenameText}
onCommitRename={commitRename}
onCancelRename={() => setRenamingId(null)}
onTogglePinned={() => void togglePinned(c)}
onDelete={() => del(c.id)}
/>
))}
@@ -931,7 +1041,12 @@ export default function ChatPage() {
// it on this render (switching to ChatView right away, before the
// async reloadConvs lands); reloadConvs then reconciles titles etc.
if (pending) setPendingByConv((p) => ({ ...p, [c.id]: pending }));
setConvs((prev) => (prev.some((x) => x.id === c.id) ? prev : [c, ...prev]));
setConvs((prev) => {
if (prev.some((item) => item.id === c.id)) return prev;
const firstUnpinned = prev.findIndex((item) => !conversationIsPinned(item));
const insertAt = firstUnpinned < 0 ? prev.length : firstUnpinned;
return [...prev.slice(0, insertAt), c, ...prev.slice(insertAt)];
});
setSelectedId(c.id);
reloadConvs();
}}
+8 -11
View File
@@ -10,7 +10,6 @@ import {
BugIcon,
ClockIcon,
NetworkIcon,
PlusIcon,
ShieldCheckIcon,
TargetIcon,
ZapIcon,
@@ -177,7 +176,7 @@ export default function DashboardPage() {
/* transient errors — next poll retries */
}
};
load();
void load();
const t = setInterval(load, 5000);
return () => {
alive = false;
@@ -213,7 +212,7 @@ export default function DashboardPage() {
/* transient errors */
}
};
load();
void load();
const t = setInterval(load, 15000);
return () => {
alive = false;
@@ -461,17 +460,11 @@ export default function DashboardPage() {
return (
<div className="flex flex-1 flex-col gap-4 pb-6">
{/* ── Header ── */}
<div className="flex items-center justify-between">
<div>
<div>
<h1 className="text-lg font-semibold tracking-tight">总览</h1>
<p className="text-xs text-muted-foreground">系统全局状态 · 实时刷新</p>
</div>
<Link href="/function/tasks">
<button className="flex items-center gap-1.5 rounded-lg bg-foreground px-3 py-1.5 text-xs font-medium text-background transition-opacity hover:opacity-90">
<PlusIcon className="size-3.5" />
新建任务
</button>
</Link>
</div>
{/* ── Row 1: 5 stat cards ── */}
@@ -575,6 +568,7 @@ export default function DashboardPage() {
] as const
).map(({ v, label }) => (
<button
type="button"
key={v}
onClick={() => setTokenVersion(v)}
title={
@@ -600,6 +594,7 @@ export default function DashboardPage() {
{/* Profile tabs */}
<div className="flex flex-wrap items-center gap-1">
<button
type="button"
onClick={() => setTokenTab("all")}
className={cn(
"rounded-md px-2.5 py-1 text-[10px] font-medium transition-colors",
@@ -616,6 +611,7 @@ export default function DashboardPage() {
const resolvedKey = tokenByProfile.has(key) ? key : p.is_default ? defaultProfileId : key;
return (
<button
type="button"
key={p.id}
onClick={() => setTokenTab(resolvedKey)}
className={cn(
@@ -745,6 +741,7 @@ export default function DashboardPage() {
] as const
).map(({ days, label }) => (
<button
type="button"
key={days}
onClick={() => setTokenDays(days)}
className={cn(
@@ -993,7 +990,7 @@ export default function DashboardPage() {
)}
</td>
<td className="px-4 py-3 tabular-nums">
{(t.in_flight ?? 0 > 0) ? (
{(t.in_flight ?? 0) > 0 ? (
<span className="font-semibold">{t.in_flight}</span>
) : (
<span className="text-muted-foreground">—</span>
File diff suppressed because it is too large. Load diff
+87 -101
View File
@@ -1,37 +1,19 @@
"use client";
import * as React from "react";
import {
TerminalIcon,
SearchIcon,
ChevronLeftIcon,
ChevronRightIcon,
Loader2Icon,
XIcon,
} from "lucide-react";
import { Card } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { ChevronLeftIcon, ChevronRightIcon, Loader2Icon, SearchIcon, TerminalIcon } from "lucide-react";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeader,
TableRow,
} from "@/components/ui/table";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { cn } from "@/lib/utils";
import { Card } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { api } from "@/lib/api";
import type { CommandRecord } from "@/lib/types";
import { cn } from "@/lib/utils";
function fmtTime(ts: string) {
return new Date(ts).toLocaleString("zh-CN", {
@@ -50,7 +32,9 @@ function toolInput(raw: string): string {
const obj = JSON.parse(raw);
if (obj && typeof obj.command === "string") return obj.command;
return JSON.stringify(obj, null, 2);
} catch { /* not JSON */ }
} catch {
/* not JSON */
}
return raw;
}
@@ -75,7 +59,7 @@ export default function CommandsPage() {
const [total, setTotal] = React.useState(0);
const [loading, setLoading] = React.useState(false);
// Inline detail panel (like traffic page, not a dialog)
// Selected execution is rendered in a right-side detail sheet.
const [selected, setSelected] = React.useState<CommandRecord | null>(null);
// Debounce search input.
@@ -85,6 +69,7 @@ export default function CommandsPage() {
}, [query]);
// Reset page on filter change.
// biome-ignore lint/correctness/useExhaustiveDependencies: every filter change intentionally resets pagination.
React.useEffect(() => {
setPage(0);
}, [queryQ, taskFilter, size]);
@@ -100,9 +85,15 @@ export default function CommandsPage() {
setCommands(r.commands ?? []);
setTotal(r.total ?? 0);
})
.catch(() => {})
.catch(() => {
if (!alive) return;
setCommands([]);
setTotal(0);
})
.finally(() => alive && setLoading(false));
return () => { alive = false; };
return () => {
alive = false;
};
}, [page, size, queryQ, taskFilter]);
const totalPages = Math.max(1, Math.ceil(total / size));
@@ -178,8 +169,8 @@ export default function CommandsPage() {
</div>
</div>
{/* History table + inline detail (Burp-style split) */}
<div className="flex h-[calc(100vh-13rem)] min-h-0 flex-col gap-3">
{/* History table */}
<div className="flex h-[calc(100vh-13rem)] min-h-0 flex-col">
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
<div className="min-h-0 flex-1 overflow-auto">
<Table>
@@ -210,18 +201,13 @@ export default function CommandsPage() {
commands.map((cmd) => (
<TableRow
key={cmd.id}
className={cn(
"cursor-pointer",
selected?.id === cmd.id && "bg-accent hover:bg-accent",
)}
className={cn("cursor-pointer", selected?.id === cmd.id && "bg-accent hover:bg-accent")}
onClick={() => setSelected(cmd)}
>
<TableCell className="text-xs text-muted-foreground tabular-nums">
{fmtTime(cmd.created_at)}
</TableCell>
<TableCell className="text-xs font-mono text-muted-foreground">
#{cmd.exploration_id}
</TableCell>
<TableCell className="text-xs font-mono text-muted-foreground">#{cmd.exploration_id}</TableCell>
<TableCell>
<Badge variant="outline" className="text-xs font-mono">
{cmd.worker || "-"}
@@ -239,9 +225,13 @@ export default function CommandsPage() {
</TableCell>
<TableCell>
{cmd.is_error ? (
<Badge variant="destructive" className="text-xs">失败</Badge>
<Badge variant="destructive" className="text-xs">
失败
</Badge>
) : (
<Badge variant="secondary" className="text-xs text-emerald-600">成功</Badge>
<Badge variant="secondary" className="text-xs text-emerald-600">
成功
</Badge>
)}
</TableCell>
</TableRow>
@@ -251,67 +241,63 @@ export default function CommandsPage() {
</Table>
</div>
</Card>
{/* Inline detail panel */}
{selected && (
<Card className="flex h-[42%] min-h-0 flex-col overflow-hidden py-0">
{/* Detail header */}
<div className="flex items-center gap-2 border-b px-3 py-2">
<Badge variant="outline" className="text-xs font-mono">
#{selected.exploration_id}
</Badge>
<Badge variant="outline" className="text-xs font-mono">
{selected.worker || "-"}
</Badge>
<Badge variant="secondary" className="text-xs font-mono">
{selected.tool || "-"}
</Badge>
<span className="text-xs text-muted-foreground">
{fmtTime(selected.created_at)}
</span>
{selected.is_error ? (
<Badge variant="destructive" className="text-xs">失败</Badge>
) : (
<Badge variant="secondary" className="text-xs text-emerald-600">成功</Badge>
)}
<Button
variant="ghost"
size="icon"
className="ml-auto size-7 shrink-0"
onClick={() => setSelected(null)}
>
<XIcon />
</Button>
</div>
{/* Input / Output split */}
<div className="grid min-h-0 flex-1 grid-cols-2 divide-x">
<div className="flex min-h-0 min-w-0 flex-col">
<div className="border-b px-3 py-1 text-[11px] font-medium text-muted-foreground">
输入 Input
</div>
<div className="min-h-0 flex-1 overflow-auto">
<pre className="p-3 font-mono text-xs break-all whitespace-pre-wrap">
{toolInput(selected.command)}
</pre>
</div>
</div>
<div className="flex min-h-0 min-w-0 flex-col">
<div className="border-b px-3 py-1 text-[11px] font-medium text-muted-foreground">
输出 Output
</div>
<div className="min-h-0 flex-1 overflow-auto">
<pre className={cn(
"p-3 font-mono text-xs break-all whitespace-pre-wrap",
selected.is_error && "text-red-600 dark:text-red-400",
)}>
{selected.output || "(空)"}
</pre>
</div>
</div>
</div>
</Card>
)}
</div>
<Sheet open={selected !== null} onOpenChange={(open) => !open && setSelected(null)}>
<SheetContent className="w-full! max-w-none! gap-0 p-0 sm:w-[46rem]! sm:max-w-[46rem]!">
{selected && (
<>
<SheetHeader className="border-b px-5 py-4">
<SheetTitle className="pr-8">工具执行详情</SheetTitle>
<SheetDescription>{fmtTime(selected.created_at)}</SheetDescription>
<div className="flex flex-wrap items-center gap-2 pt-2">
<Badge variant="outline" className="text-xs font-mono">
#{selected.exploration_id}
</Badge>
<Badge variant="outline" className="text-xs font-mono">
{selected.worker || "-"}
</Badge>
<Badge variant="secondary" className="text-xs font-mono">
{selected.tool || "-"}
</Badge>
{selected.is_error ? (
<Badge variant="destructive" className="text-xs">
失败
</Badge>
) : (
<Badge variant="secondary" className="text-xs text-emerald-600">
成功
</Badge>
)}
</div>
</SheetHeader>
<div className="grid min-h-0 flex-1 grid-rows-2 divide-y">
<div className="flex min-h-0 min-w-0 flex-col">
<div className="border-b px-5 py-2 text-xs font-medium text-muted-foreground">输入 Input</div>
<div className="min-h-0 flex-1 overflow-auto">
<pre className="p-5 font-mono text-xs break-all whitespace-pre-wrap">
{toolInput(selected.command)}
</pre>
</div>
</div>
<div className="flex min-h-0 min-w-0 flex-col">
<div className="border-b px-5 py-2 text-xs font-medium text-muted-foreground">输出 Output</div>
<div className="min-h-0 flex-1 overflow-auto">
<pre
className={cn(
"p-5 font-mono text-xs break-all whitespace-pre-wrap",
selected.is_error && "text-red-600 dark:text-red-400",
)}
>
{selected.output || "(空)"}
</pre>
</div>
</div>
</div>
</>
)}
</SheetContent>
</Sheet>
</div>
);
}
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+244 -165
View File
@@ -1,15 +1,15 @@
"use client";
import * as React from "react";
import {
RadioTowerIcon,
SearchIcon,
ChevronLeftIcon,
ChevronRightIcon,
XIcon,
Loader2Icon,
Trash2Icon,
ListChecksIcon,
Loader2Icon,
RadioTowerIcon,
SearchIcon,
Trash2Icon,
} from "lucide-react";
import {
@@ -22,35 +22,19 @@ import {
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Card } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeader,
TableRow,
} from "@/components/ui/table";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
import { Card } from "@/components/ui/card";
import { Checkbox } from "@/components/ui/checkbox";
import { cn } from "@/lib/utils";
import { Input } from "@/components/ui/input";
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { api } from "@/lib/api";
import type { TrafficExchange, TrafficResp, TrafficDetail, TrafficHost } from "@/lib/types";
import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types";
import { cn } from "@/lib/utils";
function fmtTime(ts: string) {
return new Date(ts).toLocaleString("zh-CN", {
@@ -66,7 +50,7 @@ function fmtBytes(n: number) {
if (n <= 0) return "0 B";
const units = ["B", "KB", "MB"];
const i = Math.min(Math.floor(Math.log(n) / Math.log(1024)), units.length - 1);
const v = n / Math.pow(1024, i);
const v = n / 1024 ** i;
return `${i === 0 ? v : v.toFixed(1)} ${units[i]}`;
}
@@ -78,6 +62,141 @@ function statusTone(status: number) {
return "text-muted-foreground";
}
function MethodBadge({ method }: { method: string }) {
return <Badge className="shrink-0 font-mono">{method}</Badge>;
}
// Older captures may predate Host persistence because net/http keeps Host
// outside Request.Header. Fill it for display while newly recorded traffic is
// fixed at the recorder layer as well.
function requestWithHost(raw: string, exchange: TrafficExchange): string {
if (!raw.trim() || /^host\s*:/im.test(raw)) return raw;
let host = exchange.host;
try {
host = new URL(exchange.url).host || host;
} catch {
// Relative or legacy URLs fall back to the indexed host.
}
const newline = raw.includes("\r\n") ? "\r\n" : "\n";
const firstLineEnd = raw.indexOf(newline);
if (firstLineEnd < 0) return `${raw}${newline}Host: ${host}`;
return `${raw.slice(0, firstLineEnd + newline.length)}Host: ${host}${newline}${raw.slice(firstLineEnd + newline.length)}`;
}
function StartLine({ line }: { line: string }) {
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
if (response) {
return (
<>
<span className="text-muted-foreground">{response[1]}</span>
{response[2]}
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
{response[4]}
</>
);
}
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
if (request) {
return (
<>
<span className="font-semibold text-primary">{request[1]}</span>
{request[2]}
<span className="text-chart-2">{request[3]}</span>
{request[4]}
<span className="text-muted-foreground">{request[5]}</span>
</>
);
}
return line;
}
function HeaderLine({ line }: { line: string }) {
const separator = line.indexOf(":");
if (separator <= 0) return line;
return (
<>
<span className="text-primary">{line.slice(0, separator)}</span>
<span className="text-muted-foreground">:</span>
<span className="text-chart-2">{line.slice(separator + 1)}</span>
</>
);
}
function JsonBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
let cursor = 0;
for (const match of body.matchAll(tokens)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
let className = "text-chart-4";
if (match[1]) className = match[2] ? "text-primary" : "text-chart-2";
else if (match[3]) className = "text-chart-3";
parts.push(
<span key={`${index}-${match[0].length}`} className={className}>
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
function MarkupBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tags = /<\/?[A-Za-z][^>]*>|<!--[\s\S]*?-->/g;
let cursor = 0;
for (const match of body.matchAll(tags)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
parts.push(
<span key={`${index}-${match[0].length}`} className="text-primary">
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
function HighlightedBody({ body }: { body: string }) {
const trimmed = body.trimStart();
if (trimmed.startsWith("{") || trimmed.startsWith("[")) return <JsonBody body={body} />;
if (trimmed.startsWith("<")) return <MarkupBody body={body} />;
return body;
}
function HttpCodeBlock({ raw }: { raw: string }) {
const value = raw || "(空)";
const lines = value.replaceAll("\r\n", "\n").split("\n");
const separator = lines.indexOf("");
const headerLines = separator >= 0 ? lines.slice(0, separator) : lines;
const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : "";
return (
<pre className="min-w-full whitespace-pre p-5 font-mono text-xs leading-relaxed">
<code>
{headerLines.map((line, index) => (
<React.Fragment key={`${index}-${line}`}>
{index === 0 ? <StartLine line={line} /> : <HeaderLine line={line} />}
{index < headerLines.length - 1 ? "\n" : null}
</React.Fragment>
))}
{separator >= 0 && (
<>
{"\n\n"}
<HighlightedBody body={body} />
</>
)}
</code>
</pre>
);
}
// Fixed method set (server filters exact-match); avoids deriving options from a
// single page, which would only ever list the methods on that page.
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
@@ -116,12 +235,14 @@ export default function TrafficPage() {
}, [query]);
// Any filter/size change resets to the first page.
// biome-ignore lint/correctness/useExhaustiveDependencies: these values intentionally trigger a page reset.
React.useEffect(() => {
setPage(0);
}, [hostQ, queryQ, method, size]);
// Load the current page. Auto-refresh only on page 0 (newest) so paging back
// through history isn't yanked out from under the user.
// biome-ignore lint/correctness/useExhaustiveDependencies: reloadTick is an explicit manual-refetch trigger.
React.useEffect(() => {
let alive = true;
const load = () => {
@@ -130,13 +251,17 @@ export default function TrafficPage() {
.then((r) => {
if (alive) setTraffic(r);
})
.catch(() => {});
.catch(() => {
// Keep the last successful snapshot during transient refresh failures.
});
api
.trafficHosts()
.then((r) => {
if (alive) setHosts(r.hosts ?? []);
})
.catch(() => {});
.catch(() => {
// Keep the last successful host list during transient refresh failures.
});
};
load();
const t = setInterval(() => {
@@ -155,22 +280,21 @@ export default function TrafficPage() {
const confirmDelete = () => {
setDeleting(true);
const p =
deleteMode === "selected"
? api.trafficDeleteHosts(selectedHosts)
: api.trafficDeleteHost(hostQ);
const p = deleteMode === "selected" ? api.trafficDeleteHosts(selectedHosts) : api.trafficDeleteHost(hostQ);
p.then(() => {
setDeleteMode(null);
setSelected(null);
setDetail(null);
if (deleteMode === "selected") {
setSelectedHosts([]);
setPickerOpen(false);
}
setPage(0);
setReloadTick((t) => t + 1);
setDeleteMode(null);
setSelected(null);
setDetail(null);
if (deleteMode === "selected") {
setSelectedHosts([]);
setPickerOpen(false);
}
setPage(0);
setReloadTick((t) => t + 1);
})
.catch(() => {
// Keep the confirmation open so the user can retry a failed deletion.
})
.catch(() => {})
.finally(() => setDeleting(false));
};
@@ -210,9 +334,7 @@ export default function TrafficPage() {
<div className="flex items-center justify-between gap-4">
<div>
<h1 className="text-xl font-semibold tracking-tight">流量</h1>
<p className="text-muted-foreground text-sm">
全局录制代理 · 所有 HTTP 往来
</p>
<p className="text-muted-foreground text-sm">全局录制代理 · 所有 HTTP 往来</p>
</div>
<div className="flex items-center gap-4 text-sm">
<span
@@ -226,11 +348,7 @@ export default function TrafficPage() {
<RadioTowerIcon className="size-3.5" />
{traffic?.enabled ? "录制中" : "已停用"}
</span>
{traffic?.proxy && (
<span className="font-mono text-xs text-muted-foreground">
{traffic.proxy}
</span>
)}
{traffic?.proxy && <span className="font-mono text-xs text-muted-foreground">{traffic.proxy}</span>}
<span className="text-xs text-muted-foreground">
共 <span className="tabular-nums">{traffic?.count ?? 0}</span> 条
</span>
@@ -266,29 +384,25 @@ export default function TrafficPage() {
</div>
<div className="max-h-64 overflow-y-auto">
{hosts.length === 0 ? (
<div className="px-3 py-6 text-center text-xs text-muted-foreground">
暂无流量记录
</div>
<div className="px-3 py-6 text-center text-xs text-muted-foreground">暂无流量记录</div>
) : (
hosts.map((h) => (
hosts.map((h, index) => (
<label
key={h.host}
htmlFor={`traffic-host-${index}`}
className="flex cursor-pointer items-center gap-2 px-3 py-1.5 text-xs hover:bg-accent"
>
<Checkbox
id={`traffic-host-${index}`}
checked={selectedHosts.includes(h.host)}
onCheckedChange={() =>
setSelectedHosts((prev) =>
prev.includes(h.host)
? prev.filter((x) => x !== h.host)
: [...prev, h.host],
prev.includes(h.host) ? prev.filter((x) => x !== h.host) : [...prev, h.host],
)
}
/>
<span className="truncate font-mono">{h.host}</span>
<span className="ml-auto shrink-0 tabular-nums text-muted-foreground">
{h.count}
</span>
<span className="ml-auto shrink-0 tabular-nums text-muted-foreground">{h.count}</span>
</label>
))
)}
@@ -310,12 +424,7 @@ export default function TrafficPage() {
</PopoverContent>
</Popover>
<div className="relative w-48">
<Input
placeholder="host…"
value={host}
onChange={(e) => setHost(e.target.value)}
className="h-8"
/>
<Input placeholder="host…" value={host} onChange={(e) => setHost(e.target.value)} className="h-8" />
</div>
<Button
variant="destructive"
@@ -391,8 +500,8 @@ export default function TrafficPage() {
</div>
</div>
{/* History table + detail (Burp-style split) */}
<div className="flex h-[calc(100vh-15rem)] min-h-0 flex-col gap-3">
{/* History table */}
<div className="flex h-[calc(100vh-15rem)] min-h-0 flex-col">
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
<div className="min-h-0 flex-1 overflow-auto">
<Table>
@@ -411,50 +520,29 @@ export default function TrafficPage() {
{exchanges.map((e) => (
<TableRow
key={e.id}
className={cn(
"cursor-pointer",
selected?.id === e.id && "bg-accent hover:bg-accent",
)}
className={cn("cursor-pointer", selected?.id === e.id && "bg-accent hover:bg-accent")}
onClick={() => setSelected(e)}
>
<TableCell className="text-xs text-muted-foreground tabular-nums">
{fmtTime(e.ts)}
</TableCell>
<TableCell className="text-xs text-muted-foreground tabular-nums">{fmtTime(e.ts)}</TableCell>
<TableCell className="font-mono text-xs">{e.host}</TableCell>
<TableCell>
<Badge variant="outline" className="font-mono text-xs">
{e.method}
</Badge>
<MethodBadge method={e.method} />
</TableCell>
<TableCell className="max-w-0">
<span className="block truncate font-mono text-xs">
{e.url}
</span>
<span className="block truncate font-mono text-xs">{e.url}</span>
</TableCell>
<TableCell>
<span
className={cn(
"font-mono text-xs font-semibold tabular-nums",
statusTone(e.status),
)}
>
<span className={cn("font-mono text-xs font-semibold tabular-nums", statusTone(e.status))}>
{e.status}
</span>
</TableCell>
<TableCell className="text-xs text-muted-foreground">
{e.content_type}
</TableCell>
<TableCell className="text-right text-xs tabular-nums">
{fmtBytes(e.resp_len)}
</TableCell>
<TableCell className="text-xs text-muted-foreground">{e.content_type}</TableCell>
<TableCell className="text-right text-xs tabular-nums">{fmtBytes(e.resp_len)}</TableCell>
</TableRow>
))}
{exchanges.length === 0 && (
<TableRow>
<TableCell
colSpan={7}
className="py-12 text-center text-sm text-muted-foreground"
>
<TableCell colSpan={7} className="py-12 text-center text-sm text-muted-foreground">
{traffic === null ? "加载中…" : "没有匹配的流量。"}
</TableCell>
</TableRow>
@@ -463,66 +551,60 @@ export default function TrafficPage() {
</Table>
</div>
</Card>
{/* Raw request/response pane */}
{selected && (
<Card className="flex h-[42%] min-h-0 flex-col overflow-hidden py-0">
<div className="flex items-center gap-2 border-b px-3 py-2">
<Badge variant="outline" className="font-mono text-xs">
{selected.method}
</Badge>
<span className="truncate font-mono text-xs">
{selected.host}
{selected.url}
</span>
<span
className={cn(
"ml-auto font-mono text-xs font-semibold tabular-nums",
statusTone(selected.status),
)}
>
{selected.status}
</span>
<Button
variant="ghost"
size="icon"
className="size-7 shrink-0"
onClick={() => setSelected(null)}
>
<XIcon />
</Button>
</div>
<div className="grid min-h-0 flex-1 grid-cols-2 divide-x">
{(
[
["请求 Request", detail?.req],
["响应 Response", detail?.resp],
] as const
).map(([label, body]) => (
<div key={label} className="flex min-h-0 min-w-0 flex-col">
<div className="border-b px-3 py-1 text-[11px] font-medium text-muted-foreground">
{label}
</div>
<div className="min-h-0 flex-1 overflow-auto">
{detailLoading ? (
<div className="flex items-center gap-2 p-3 text-xs text-muted-foreground">
<Loader2Icon className="size-3.5 animate-spin" />
加载报文…
</div>
) : (
<pre className="p-3 font-mono text-xs break-all whitespace-pre-wrap">
{body || "(空)"}
</pre>
)}
</div>
</div>
))}
</div>
</Card>
)}
</div>
<AlertDialog open={deleteMode !== null} onOpenChange={(o) => { if (!o) setDeleteMode(null); }}>
<Sheet open={selected !== null} onOpenChange={(open) => !open && setSelected(null)}>
<SheetContent className="w-full! max-w-none! gap-0 p-0 sm:w-[48rem]! sm:max-w-[48rem]!">
{selected && (
<>
<SheetHeader className="border-b px-5 py-4">
<div className="flex items-center gap-2 pr-8">
<MethodBadge method={selected.method} />
<Badge variant="secondary" className={cn("font-mono tabular-nums", statusTone(selected.status))}>
{selected.status}
</Badge>
<span className="ml-auto text-xs text-muted-foreground tabular-nums">{fmtTime(selected.ts)}</span>
</div>
<SheetTitle className="break-all font-mono">{selected.host}</SheetTitle>
<SheetDescription className="break-all font-mono">{selected.url}</SheetDescription>
</SheetHeader>
<Tabs defaultValue="request" className="min-h-0 flex-1 gap-0">
<TabsList className="mx-5 mt-4 grid w-auto grid-cols-2">
<TabsTrigger value="request">请求 Request</TabsTrigger>
<TabsTrigger value="response">响应 Response</TabsTrigger>
</TabsList>
<TabsContent value="request" className="min-h-0 overflow-auto">
{detailLoading ? (
<div className="flex items-center gap-2 p-5 text-xs text-muted-foreground">
<Loader2Icon className="size-3.5 animate-spin" />
加载报文…
</div>
) : (
<HttpCodeBlock raw={requestWithHost(detail?.req ?? "", selected)} />
)}
</TabsContent>
<TabsContent value="response" className="min-h-0 overflow-auto">
{detailLoading ? (
<div className="flex items-center gap-2 p-5 text-xs text-muted-foreground">
<Loader2Icon className="size-3.5 animate-spin" />
加载报文…
</div>
) : (
<HttpCodeBlock raw={detail?.resp ?? ""} />
)}
</TabsContent>
</Tabs>
</>
)}
</SheetContent>
</Sheet>
<AlertDialog
open={deleteMode !== null}
onOpenChange={(o) => {
if (!o) setDeleteMode(null);
}}
>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>
@@ -533,9 +615,7 @@ export default function TrafficPage() {
<AlertDialogDescription>
{deleteMode === "selected" ? (
<>
将永久删除{" "}
<span className="font-semibold tabular-nums">{selectedHosts.length}</span>{" "}
个目标(
将永久删除 <span className="font-semibold tabular-nums">{selectedHosts.length}</span> 个目标(
<span className="font-mono">
{selectedHosts.slice(0, 3).join("、")}
{selectedHosts.length > 3 ? "…" : ""}
@@ -544,8 +624,7 @@ export default function TrafficPage() {
</>
) : (
<>
将永久删除 host 包含{" "}
<span className="font-mono font-semibold">{hostQ}</span>{" "}
将永久删除 host 包含 <span className="font-mono font-semibold">{hostQ}</span>{" "}
的所有流量记录(含请求/响应原文),此操作不可撤销。
</>
)}
+21 -30
View File
@@ -1,22 +1,12 @@
"use client";
import * as React from "react";
import { useEffect, useRef } from "react";
import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
Pagination,
PaginationContent,
PaginationEllipsis,
PaginationItem,
} from "@/components/ui/pagination";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { Pagination, PaginationContent, PaginationEllipsis, PaginationItem } from "@/components/ui/pagination";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
interface TablePaginationProps {
page: number; // 1-based
@@ -49,19 +39,20 @@ export function TablePagination({
pageSizeOptions = [10, 20, 50],
}: TablePaginationProps) {
const totalPages = Math.max(1, Math.ceil(total / pageSize));
const from = total === 0 ? 0 : (page - 1) * pageSize + 1;
const to = Math.min(page * pageSize, total);
const safePage = Math.min(Math.max(1, page), totalPages);
const from = total === 0 ? 0 : (safePage - 1) * pageSize + 1;
const to = Math.min(safePage * pageSize, total);
const onPageChangeRef = useRef(onPageChange);
onPageChangeRef.current = onPageChange;
useEffect(() => {
if (page !== safePage) onPageChangeRef.current(safePage);
}, [page, safePage]);
return (
<div className="flex flex-wrap items-center justify-between gap-2 border-t px-4 py-2 lg:px-6">
<div className="text-muted-foreground flex items-center gap-2 text-xs">
<Select
value={String(pageSize)}
onValueChange={(v) => {
onPageSizeChange(Number(v));
onPageChange(1);
}}
>
<Select value={String(pageSize)} onValueChange={(v) => onPageSizeChange(Number(v))}>
<SelectTrigger size="sm" className="h-7 w-16">
<SelectValue />
</SelectTrigger>
@@ -90,14 +81,14 @@ export function TablePagination({
<Button
variant="ghost"
size="icon-sm"
disabled={page === 1}
onClick={() => onPageChange(page - 1)}
disabled={safePage === 1}
onClick={() => onPageChange(safePage - 1)}
aria-label="上一页"
>
<ChevronLeftIcon className="size-4" />
</Button>
</PaginationItem>
{pageWindows(page, totalPages).map((p, i) =>
{pageWindows(safePage, totalPages).map((p, i) =>
p === "..." ? (
<PaginationItem key={`el-${i}`}>
<PaginationEllipsis />
@@ -105,10 +96,10 @@ export function TablePagination({
) : (
<PaginationItem key={p}>
<Button
variant={p === page ? "outline" : "ghost"}
variant={p === safePage ? "outline" : "ghost"}
size="icon-sm"
onClick={() => onPageChange(p)}
aria-current={p === page ? "page" : undefined}
aria-current={p === safePage ? "page" : undefined}
>
{p}
</Button>
@@ -119,8 +110,8 @@ export function TablePagination({
<Button
variant="ghost"
size="icon-sm"
disabled={page === totalPages}
onClick={() => onPageChange(page + 1)}
disabled={safePage === totalPages}
onClick={() => onPageChange(safePage + 1)}
aria-label="下一页"
>
<ChevronRightIcon className="size-4" />
@@ -0,0 +1,459 @@
"use client";
import * as React from "react";
import { LibraryIcon, PlusIcon, SaveIcon, Settings2Icon, Trash2Icon } from "lucide-react";
import { toast } from "sonner";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Sheet, SheetContent, SheetDescription, SheetFooter, SheetHeader, SheetTitle } from "@/components/ui/sheet";
import { Spinner } from "@/components/ui/spinner";
import { Textarea } from "@/components/ui/textarea";
import { api } from "@/lib/api";
import type { TaskTemplate } from "@/lib/types";
import { cn } from "@/lib/utils";
interface TemplateDraft {
name: string;
description: string;
goal: string;
}
interface TaskTemplateManagerProps {
open: boolean;
onOpenChange: (open: boolean) => void;
templates: TaskTemplate[];
seed: Pick<TemplateDraft, "description" | "goal"> | null;
onCreated: (template: TaskTemplate) => void;
onUpdated: (template: TaskTemplate) => void;
onDeleted: (id: number) => void;
}
const emptyDraft = (): TemplateDraft => ({ name: "", description: "", goal: "" });
function templateDraft(template: TaskTemplate): TemplateDraft {
return {
name: template.name,
description: template.description,
goal: template.goal,
};
}
function TaskTemplateManager({
open,
onOpenChange,
templates,
seed,
onCreated,
onUpdated,
onDeleted,
}: TaskTemplateManagerProps) {
const [selectedID, setSelectedID] = React.useState<number | null>(null);
const [draft, setDraft] = React.useState<TemplateDraft>(emptyDraft);
const [saving, setSaving] = React.useState(false);
const [deleteOpen, setDeleteOpen] = React.useState(false);
const [deleting, setDeleting] = React.useState(false);
const wasOpen = React.useRef(false);
React.useEffect(() => {
if (open && !wasOpen.current) {
if (seed) {
setSelectedID(null);
setDraft({ name: "", description: seed.description, goal: seed.goal });
} else if (templates[0]) {
setSelectedID(templates[0].id);
setDraft(templateDraft(templates[0]));
} else {
setSelectedID(null);
setDraft(emptyDraft());
}
}
wasOpen.current = open;
}, [open, seed, templates]);
const selectTemplate = (template: TaskTemplate) => {
setSelectedID(template.id);
setDraft(templateDraft(template));
};
const startNew = () => {
setSelectedID(null);
setDraft(emptyDraft());
};
const updateDraft = (field: keyof TemplateDraft, value: string) => {
setDraft((current) => ({ ...current, [field]: value }));
};
async function save() {
const input = {
name: draft.name.trim(),
description: draft.description.trim(),
goal: draft.goal.trim(),
};
if (!input.name || !input.description || !input.goal) {
toast.error("请填写模板名称、描述和目标");
return;
}
setSaving(true);
try {
if (selectedID == null) {
const created = await api.createTaskTemplate(input);
onCreated(created);
setSelectedID(created.id);
setDraft(templateDraft(created));
toast.success("模板已创建");
} else {
const updated = await api.updateTaskTemplate(selectedID, input);
onUpdated(updated);
setDraft(templateDraft(updated));
toast.success("模板已更新");
}
} catch (error) {
toast.error(`保存失败:${(error as Error).message}`);
} finally {
setSaving(false);
}
}
async function remove() {
if (selectedID == null) return;
const deletedID = selectedID;
setDeleting(true);
try {
await api.deleteTaskTemplate(deletedID);
onDeleted(deletedID);
const next = templates.find((template) => template.id !== deletedID);
if (next) {
setSelectedID(next.id);
setDraft(templateDraft(next));
} else {
startNew();
}
setDeleteOpen(false);
toast.success("模板已删除");
} catch (error) {
toast.error(`删除失败:${(error as Error).message}`);
} finally {
setDeleting(false);
}
}
let saveLabel = saving ? "保存中" : "保存修改";
if (!saving && selectedID == null) saveLabel = "创建模板";
return (
<>
<Sheet open={open} onOpenChange={onOpenChange}>
<SheetContent className="grid h-full w-full! max-w-none! grid-rows-[auto_minmax(0,1fr)_auto] gap-0 overflow-hidden p-0 sm:w-[48rem]! sm:max-w-[48rem]!">
<SheetHeader className="border-b px-6 py-5">
<SheetTitle>任务模板管理</SheetTitle>
<SheetDescription>模板只保存描述与目标;修改不会影响已经创建的任务。</SheetDescription>
</SheetHeader>
<div className="grid min-h-0 overflow-y-auto lg:grid-cols-[15rem_minmax(0,1fr)] lg:overflow-hidden">
<div className="flex min-h-0 flex-col border-b p-3 lg:border-r lg:border-b-0">
<Button type="button" variant="outline" className="w-full" onClick={startNew}>
<PlusIcon data-icon="inline-start" />
新建模板
</Button>
<ScrollArea className="mt-2 max-h-44 lg:max-h-none lg:flex-1">
<div className="flex flex-col gap-1 pr-2">
{templates.length === 0 && (
<p className="px-2 py-6 text-center text-muted-foreground text-sm">暂无模板</p>
)}
{templates.map((template) => (
<button
key={template.id}
type="button"
className={cn(
"min-w-0 rounded-md px-2.5 py-2 text-left transition-colors",
selectedID === template.id ? "bg-accent text-accent-foreground" : "hover:bg-accent/50",
)}
onClick={() => selectTemplate(template)}
>
<span className="block truncate font-medium text-sm">{template.name}</span>
<span className="block truncate text-muted-foreground text-xs">{template.description}</span>
</button>
))}
</div>
</ScrollArea>
</div>
<ScrollArea className="min-h-0">
<FieldGroup className="p-6">
<Field>
<FieldLabel htmlFor="task-template-name">模板名称</FieldLabel>
<Input
id="task-template-name"
value={draft.name}
maxLength={120}
placeholder="例如:外部 Web 渗透"
onChange={(event) => updateDraft("name", event.target.value)}
/>
</Field>
<Field>
<FieldLabel htmlFor="task-template-description">描述</FieldLabel>
<Textarea
id="task-template-description"
className="min-h-28"
value={draft.description}
placeholder="测试对象与背景"
onChange={(event) => updateDraft("description", event.target.value)}
/>
</Field>
<Field>
<FieldLabel htmlFor="task-template-goal">目标</FieldLabel>
<Textarea
id="task-template-goal"
className="min-h-28"
value={draft.goal}
placeholder="任务需要达成的目标"
onChange={(event) => updateDraft("goal", event.target.value)}
/>
</Field>
</FieldGroup>
</ScrollArea>
</div>
<SheetFooter className="border-t px-6 py-4 sm:flex-row sm:items-center">
{selectedID != null && (
<Button type="button" variant="destructive" className="sm:mr-auto" onClick={() => setDeleteOpen(true)}>
<Trash2Icon data-icon="inline-start" />
删除模板
</Button>
)}
<Button type="button" variant="outline" onClick={() => onOpenChange(false)}>
关闭
</Button>
<Button type="button" disabled={saving} onClick={() => void save()}>
{saving ? <Spinner data-icon="inline-start" /> : <SaveIcon data-icon="inline-start" />}
{saveLabel}
</Button>
</SheetFooter>
</SheetContent>
</Sheet>
<AlertDialog open={deleteOpen} onOpenChange={setDeleteOpen}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>删除模板「{draft.name || "未命名模板"}」?</AlertDialogTitle>
<AlertDialogDescription>已由该模板创建的任务不会受到影响。</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={deleting}>取消</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
disabled={deleting}
onClick={(event) => {
event.preventDefault();
void remove();
}}
>
{deleting && <Spinner data-icon="inline-start" />}
{deleting ? "删除中" : "删除"}
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</>
);
}
interface TaskTemplateControlsProps {
description: string;
goal: string;
selectedTemplateID: number | null;
onSelectedTemplateIDChange: (id: number | null) => void;
onApply: (template: TaskTemplate) => void;
portalContainer?: React.RefObject<HTMLElement | null>;
}
export function TaskTemplateControls({
description,
goal,
selectedTemplateID,
onSelectedTemplateIDChange,
onApply,
portalContainer,
}: TaskTemplateControlsProps) {
const [templates, setTemplates] = React.useState<TaskTemplate[]>([]);
const [loading, setLoading] = React.useState(true);
const [templateInputValue, setTemplateInputValue] = React.useState("");
const [pendingTemplate, setPendingTemplate] = React.useState<TaskTemplate | null>(null);
const [managerOpen, setManagerOpen] = React.useState(false);
const [managerSeed, setManagerSeed] = React.useState<Pick<TemplateDraft, "description" | "goal"> | null>(null);
const loadTemplates = React.useCallback(async () => {
setLoading(true);
try {
setTemplates(await api.taskTemplates());
} catch (error) {
setTemplates([]);
toast.error(`加载模板失败:${(error as Error).message}`);
} finally {
setLoading(false);
}
}, []);
React.useEffect(() => {
void loadTemplates();
}, [loadTemplates]);
const selectedTemplate = React.useMemo(
() => templates.find((template) => template.id === selectedTemplateID) ?? null,
[selectedTemplateID, templates],
);
React.useEffect(() => {
setTemplateInputValue(selectedTemplate?.name ?? "");
}, [selectedTemplate]);
const applyTemplate = (template: TaskTemplate) => {
onApply(template);
onSelectedTemplateIDChange(template.id);
setPendingTemplate(null);
};
const chooseTemplate = (template: TaskTemplate | null) => {
if (!template) {
setTemplateInputValue("");
onSelectedTemplateIDChange(null);
return;
}
setTemplateInputValue(template.name);
const hasContent = description.trim() !== "" || goal.trim() !== "";
const changesContent = description !== template.description || goal !== template.goal;
if (hasContent && changesContent) {
setPendingTemplate(template);
return;
}
applyTemplate(template);
};
const openManager = (seed: Pick<TemplateDraft, "description" | "goal"> | null) => {
setManagerSeed(seed);
setManagerOpen(true);
};
const upsertTemplate = (template: TaskTemplate) => {
setTemplates((current) => [template, ...current.filter((item) => item.id !== template.id)]);
};
const deleteTemplate = (id: number) => {
setTemplates((current) => current.filter((template) => template.id !== id));
if (selectedTemplateID === id) onSelectedTemplateIDChange(null);
};
let pickerPlaceholder = loading ? "正在加载模板" : "暂无任务模板";
if (!loading && templates.length > 0) pickerPlaceholder = "搜索并选择任务模板";
return (
<>
<Field>
<div className="flex flex-wrap items-center justify-between gap-2">
<FieldLabel htmlFor="task-template-picker">任务模板</FieldLabel>
<div className="flex flex-wrap gap-2">
<Button type="button" variant="ghost" size="sm" onClick={() => openManager(null)}>
<Settings2Icon data-icon="inline-start" />
管理模板
</Button>
<Button
type="button"
variant="ghost"
size="sm"
disabled={!description.trim() || !goal.trim()}
onClick={() => openManager({ description, goal })}
>
<SaveIcon data-icon="inline-start" />
另存为模板
</Button>
</div>
</div>
<Combobox
items={templates}
itemToStringLabel={(template) => template.name}
itemToStringValue={(template) => String(template.id)}
inputValue={templateInputValue}
onInputValueChange={(value) => setTemplateInputValue(value)}
value={selectedTemplate}
onValueChange={chooseTemplate}
>
<ComboboxInput
id="task-template-picker"
className="w-full"
placeholder={pickerPlaceholder}
disabled={loading || templates.length === 0}
showClear
/>
<ComboboxContent portalContainer={portalContainer}>
<ComboboxEmpty>没有匹配的模板</ComboboxEmpty>
<ComboboxList>
{(template) => (
<ComboboxItem key={template.id} value={template}>
<LibraryIcon />
<div className="min-w-0 flex-1">
<p className="truncate font-medium text-sm">{template.name}</p>
{template.description && (
<p className="truncate text-muted-foreground text-xs">{template.description}</p>
)}
</div>
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
<FieldDescription>选择后会复制模板的描述和目标,不与模板保持关联。</FieldDescription>
</Field>
<AlertDialog
open={pendingTemplate != null}
onOpenChange={(open) => {
if (open) return;
setPendingTemplate(null);
setTemplateInputValue(selectedTemplate?.name ?? "");
}}
>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>使用模板「{pendingTemplate?.name}」?</AlertDialogTitle>
<AlertDialogDescription>当前已填写的描述和目标将被模板内容覆盖。</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>取消</AlertDialogCancel>
<AlertDialogAction onClick={() => pendingTemplate && applyTemplate(pendingTemplate)}>
覆盖并使用
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
<TaskTemplateManager
open={managerOpen}
onOpenChange={setManagerOpen}
templates={templates}
seed={managerSeed}
onCreated={upsertTemplate}
onUpdated={upsertTemplate}
onDeleted={deleteTemplate}
/>
</>
);
}
+3 -3
View File
@@ -4,7 +4,7 @@ import { cva, type VariantProps } from "class-variance-authority"
import { cn } from "@/lib/utils"
const alertVariants = cva(
"group/alert relative grid w-full gap-0.5 rounded-lg border px-2.5 py-2 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4",
"group/alert relative grid w-full gap-0.5 rounded-lg border px-2.5 py-2 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_minmax(0,1fr)] has-[>svg]:gap-x-2 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4",
{
variants: {
variant: {
@@ -39,7 +39,7 @@ function AlertTitle({ className, ...props }: React.ComponentProps<"div">) {
<div
data-slot="alert-title"
className={cn(
"font-medium group-has-[>svg]/alert:col-start-2 [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground",
"min-w-0 font-medium group-has-[>svg]/alert:col-start-2 [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground",
className
)}
{...props}
@@ -55,7 +55,7 @@ function AlertDescription({
<div
data-slot="alert-description"
className={cn(
"text-sm text-balance text-muted-foreground md:text-pretty [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground [&_p:not(:last-child)]:mb-4",
"min-w-0 text-sm text-balance text-muted-foreground [overflow-wrap:anywhere] md:text-pretty [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground [&_p:not(:last-child)]:mb-4",
className
)}
{...props}
+46 -28
View File
@@ -11,11 +11,12 @@ import type {
AgentDetail,
AgentTrigger,
Asset,
AssetNode,
Audit,
BatchControlItem,
ChatAttachment,
CommandRecord,
Company,
CompanyScopeRule,
Conversation,
ConvTokenSummary,
CoverageAssetRefs,
@@ -25,6 +26,8 @@ import type {
DeleteTaskResult,
Edge,
Finding,
FindingDeepenResponse,
FindingGroupsPage,
FindingQuery,
FindingStats,
FindingStatus,
@@ -39,21 +42,23 @@ import type {
LLMTask,
MCPServer,
MCPTool,
MissingSkill,
ModelTokenStat,
PromptVar,
PromptVersion,
SessionTokenUsage,
Settings,
Severity,
SkillCall,
SkillItem,
MissingSkill,
SSProject,
SSTask,
Stats,
Task,
TaskAssetView,
TaskLLMResolutions,
TaskNode,
TaskScopeRow,
TaskTemplate,
TokenTotal,
TokenUsage,
Tool,
@@ -179,6 +184,7 @@ export const api = {
goal: string;
llmProfileIds?: number[];
sourceTaskIds?: string[];
companyIds?: number[];
timeoutSeconds?: number;
seedFirstIntent?: boolean;
planHeartbeatSeconds?: number;
@@ -188,10 +194,17 @@ export const api = {
goal: input.goal,
llm_profile_ids: input.llmProfileIds ?? [],
source_task_ids: input.sourceTaskIds ?? [],
company_ids: input.companyIds ?? [],
timeout_seconds: input.timeoutSeconds ?? 0,
seed_first_intent: input.seedFirstIntent ?? false,
plan_heartbeat_seconds: input.planHeartbeatSeconds ?? 0, // 0 = 后端归一到默认 600(10min)
}),
taskTemplates: () => get<{ templates: TaskTemplate[] }>("/task-templates").then((r) => arr(r.templates)),
createTaskTemplate: (input: Pick<TaskTemplate, "name" | "description" | "goal">) =>
post<TaskTemplate>("/task-templates", input),
updateTaskTemplate: (id: number, input: Partial<Pick<TaskTemplate, "name" | "description" | "goal">>) =>
patch<TaskTemplate>(`/task-templates/${id}`, input),
deleteTaskTemplate: (id: number) => del<{ deleted: number }>(`/task-templates/${id}`),
updateTaskLLMProfiles: (id: string, llmProfileIds: number[], activeLLMProfileId?: number) =>
put<{
id: string;
@@ -206,13 +219,16 @@ export const api = {
}),
deleteTask: (id: string, options: DeleteTaskOptions) => del<DeleteTaskResult>(`/tasks/${id}`, options),
controlTask: (id: string, action: "pause" | "resume") =>
post<{ id: string; paused: boolean }>(`/tasks/${id}/control`, { action }),
post<{ id: string; paused: boolean; queued: boolean; status: string }>(`/tasks/${id}/control`, { action }),
controlTasksBatch: (taskIds: string[], action: "pause" | "resume") =>
post<{ items: BatchControlItem[] }>("/tasks/control/batch", { task_ids: taskIds, action }),
controlIntent: (taskId: string, intentId: string, action: "pause" | "resume" | "cancel") =>
post<{
id: number;
state: "paused" | "open" | "cancelled";
deleted?: { intents: number; facts: number; findings: number; activities: number };
}>(`/tasks/${taskId}/intents/${intentId}/control`, { action }),
taskLLMResolution: (id: string) => get<TaskLLMResolutions>(`/tasks/${id}/llm/resolution`),
// 重跑一条没跑成功的意图(blocked/exhausted/stopped):置回 open,worker 会重新认领、从头再跑。
rerunIntent: (taskId: string, intentId: string) =>
post<{ id: string; reopened: number }>(`/tasks/${taskId}/intents/${intentId}/rerun`),
@@ -308,38 +324,22 @@ export const api = {
// ---- companies (企业 + 资产范围;归属唯一来源) ----
companies: () => get<Company[]>("/companies").then(arr),
createCompany: (name: string, logo: string, scope: string) =>
createCompany: (name: string, scope: CompanyScopeRule[]) =>
post<{ id: number; created: boolean; scope_added?: number; scope_invalid?: number; scope_errors?: string[] }>(
"/companies",
{
name,
logo,
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
},
),
addCompanyScope: (id: number, scope: string, reason = "") =>
addCompanyScope: (id: number, scope: CompanyScopeRule[], reason = "") =>
post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, {
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
reason,
}),
updateCompanyScope: (id: number, scope: string, reason = "") =>
updateCompanyScope: (id: number, scope: CompanyScopeRule[], reason = "") =>
post<{ added: number; skipped: number; invalid: number; errors?: string[] }>(`/companies/${id}/scope`, {
scope: scope.trim()
? scope
.split("\n")
.map((s) => s.trim())
.filter(Boolean)
: [],
scope,
reason,
reset: true,
}),
@@ -361,6 +361,15 @@ export const api = {
if (q.sort) p.set("sort", q.sort);
return get<FindingsPage>(`/exploration/findings?${p.toString()}`);
},
findingGroups: (q: FindingQuery) => {
const p = new URLSearchParams({ page: String(q.page), limit: String(q.pageSize) });
if (q.severity && q.severity !== "all") p.set("severity", q.severity);
if (q.status && q.status !== "all") p.set("status", q.status);
if (q.vulnclass && q.vulnclass !== "all") p.set("vulnclass", q.vulnclass);
if (q.task && q.task !== "all") p.set("task_id", q.task);
if (q.sort) p.set("sort", q.sort);
return get<FindingGroupsPage>(`/exploration/findings/groups?${p.toString()}`);
},
findingStats: () => get<FindingStats>("/exploration/findings/stats"),
// exportFindings 触发发现页导出并下载文件。scope=selected 时传 ids(finding_id 列表);
// scope=filtered 时传当前筛选(沿用 FindingQuery 的筛选字段);scope=all 忽略筛选。
@@ -415,10 +424,15 @@ export const api = {
) => patch<Finding>(`/exploration/findings/${id}`, fields),
// 删除漏洞:移除 findings 记录 + 来源探索节点(从发现列表/任务发现 Tab/探索图一并消失)。
deleteFinding: (id: string) => del<{ deleted: boolean; id: number }>(`/exploration/findings/${id}`),
deepenFinding: (id: string, description: string) =>
post<FindingDeepenResponse>(`/exploration/findings/${id}/deepen`, { description }),
intents: (task?: string) => get<TaskNode[]>(`/exploration/intents${tq(task)}`).then(arr),
tokenStats: (task?: string) =>
get<{ workers: TokenUsage[]; total: TokenTotal }>(`/exploration/tokens${tq(task)}`).then((r) => ({
get<{ workers: TokenUsage[]; sessions: SessionTokenUsage[]; total: TokenTotal }>(
`/exploration/tokens${tq(task)}`,
).then((r) => ({
workers: arr(r.workers),
sessions: arr(r.sessions),
total: r.total,
})),
tokenDaily: (days = 30) => get<DailyTokenBucket[]>(`/tokens/daily?days=${days}`).then(arr),
@@ -505,6 +519,7 @@ export const api = {
},
chat: (message: string, task?: string, attachments?: ChatAttachment[]) =>
post<{ reply: string; mode: string }>(`/chat${tq(task)}`, { message, attachments }),
chatStatus: (taskId: string) => get<{ running: boolean }>(`/tasks/${taskId}/chat/status`),
// 方式1 文件上传:落到会话/任务工作目录 uploads/,返回可供 agent Read 的相对路径。
chatUpload: async (scope: "task" | "session" | "staging", id: string, files: File[]) => {
if (MOCK)
@@ -619,7 +634,10 @@ export const api = {
conversations: () => get<{ conversations: Conversation[] }>("/conversations").then((r) => arr(r.conversations)),
createConversation: (agent_key: string, title = "", llm_profile_id?: number | null) =>
post<Conversation>("/conversations", { agent_key, title, llm_profile_id: llm_profile_id ?? null }),
renameConversation: (id: number, title: string) => patch<{ ok: boolean }>(`/conversations/${id}`, { title }),
updateConversation: (id: number, input: { title?: string; pinned?: boolean }) =>
patch<Conversation>(`/conversations/${id}`, input),
renameConversation: (id: number, title: string) => patch<Conversation>(`/conversations/${id}`, { title }),
pinConversation: (id: number, pinned: boolean) => patch<Conversation>(`/conversations/${id}`, { pinned }),
updateConversationProfile: (id: number, llm_profile_id: number | null) =>
patch<{ ok: boolean }>(`/conversations/${id}/profile`, { llm_profile_id }),
deleteConversation: (id: number) => del<{ deleted: number }>(`/conversations/${id}`),
+212
View File
@@ -0,0 +1,212 @@
import type { CompanyScopeKind, CompanyScopeRule } from "@/lib/types";
export const MAX_COMPANY_SCOPE_RULES = 256;
export const MAX_COMPANY_SCOPE_VALUE_LENGTH = 1024;
export type CompanyScopeTextIssue = {
line: number;
rule?: CompanyScopeRule;
error?: string;
};
export type ParsedCompanyScopeText = {
rules: CompanyScopeRule[];
errors: Array<{ line: number; error: string }>;
};
export const COMPANY_SCOPE_KINDS = new Set<CompanyScopeKind>(["domain", "ip", "cidr", "icp", "keyword"]);
export function isCompanyScopeKind(value: unknown): value is CompanyScopeKind {
return typeof value === "string" && COMPANY_SCOPE_KINDS.has(value as CompanyScopeKind);
}
function isIPv4(value: string): boolean {
const parts = value.split(".");
return (
parts.length === 4 && parts.every((part) => (part === "0" || /^[1-9]\d{0,2}$/.test(part)) && Number(part) <= 255)
);
}
function isIPv6(value: string): boolean {
if (!value.includes(":") || /[^0-9a-f:.]/i.test(value) || value.includes(":::")) return false;
const halves = value.split("::");
if (halves.length > 2) return false;
const countSegments = (half: string): number | null => {
if (!half) return 0;
const segments = half.split(":");
let count = 0;
for (const [index, segment] of segments.entries()) {
if (segment.includes(".")) {
if (index !== segments.length - 1 || !isIPv4(segment)) return null;
count += 2;
} else {
if (!/^[0-9a-f]{1,4}$/i.test(segment)) return null;
count++;
}
}
return count;
};
const left = countSegments(halves[0]);
const right = countSegments(halves[1] ?? "");
if (left === null || right === null) return false;
return halves.length === 2 ? left + right < 8 : left === 8;
}
function ipVersion(value: string): 4 | 6 | null {
if (isIPv4(value)) return 4;
if (isIPv6(value)) return 6;
return null;
}
function looksLikeIPAddress(value: string): boolean {
const trimmed = value.trim();
if (trimmed.includes("://") || /\s/.test(trimmed)) return false;
if ((trimmed.match(/:/g) ?? []).length >= 2) {
const parts = trimmed.split(":");
let validSegments = 0;
for (const part of parts) {
if (!part) {
if (validSegments > 0 || trimmed.startsWith("::")) return true;
continue;
}
if (part.length > 4 || !/^[0-9a-f]+$/i.test(part)) return false;
validSegments++;
if (validSegments >= 2) return true;
}
return false;
}
return trimmed.includes(".") && /^[0-9.]+$/.test(trimmed);
}
function domainHostname(value: string): string | null {
try {
let candidate = value;
if (value.startsWith("//")) candidate = `http:${value}`;
else if (!value.includes("://")) candidate = `http://${value}`;
const url = new URL(candidate);
return (
url.hostname
.replace(/^\[|\]$/g, "")
.replace(/\.$/, "")
.toLowerCase() || null
);
} catch {
return null;
}
}
export function companyScopeRuleError(rule: CompanyScopeRule): string {
const value = rule.value.trim();
if (!value) return "请填写范围值";
if (Array.from(value).length > MAX_COMPANY_SCOPE_VALUE_LENGTH) {
return `最多 ${MAX_COMPANY_SCOPE_VALUE_LENGTH} 个字符`;
}
if (rule.kind === "domain") {
if (/\s/.test(value)) return "请输入不含空格的有效域名或 URL";
const hostname = domainHostname(value);
if (!hostname || ipVersion(hostname) !== null) return "请输入有效域名或 URL";
const labels = hostname.split(".");
if (
hostname.length > 253 ||
labels.length < 2 ||
labels.some(
(label) =>
!label || label.length > 63 || label.startsWith("-") || label.endsWith("-") || !/^[a-z0-9-]+$/i.test(label),
)
) {
return "请输入有效域名或 URL";
}
}
if (rule.kind === "ip" && ipVersion(value) === null) return "请输入有效 IP";
if (rule.kind === "cidr") {
const separator = value.lastIndexOf("/");
if (separator <= 0) return "请输入 CIDR 网段";
const address = value.slice(0, separator);
const prefixText = value.slice(separator + 1);
const version = ipVersion(address);
if (version === null || !/^\d+$/.test(prefixText)) return "请输入 CIDR 网段";
const prefix = Number(prefixText);
if (version === 4 && (prefix < 16 || prefix > 32)) return "IPv4 网段前缀需为 /16 至 /32";
if (version === 6 && (prefix < 32 || prefix > 128)) return "IPv6 网段前缀需为 /32 至 /128";
}
return "";
}
export function classifyCompanyScopeLine(
raw: string,
line: number,
preservedRule?: CompanyScopeRule,
): CompanyScopeTextIssue {
const value = raw.trim();
if (Array.from(value).length > MAX_COMPANY_SCOPE_VALUE_LENGTH) {
return { line, error: `最多 ${MAX_COMPANY_SCOPE_VALUE_LENGTH} 个字符` };
}
if (preservedRule) {
const rule = { kind: preservedRule.kind, value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
const separator = value.lastIndexOf("/");
if (
separator > 0 &&
(ipVersion(value.slice(0, separator)) !== null || looksLikeIPAddress(value.slice(0, separator)))
) {
const rule: CompanyScopeRule = { kind: "cidr", value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
if (ipVersion(value) !== null) return { line, rule: { kind: "ip", value } };
if (looksLikeIPAddress(value)) return { line, error: "请输入有效 IP" };
const looksLikeDomain = value.includes("://") || (!/\s/.test(value) && value.includes("."));
if (looksLikeDomain) {
const hostname = domainHostname(value);
if (hostname && ipVersion(hostname) !== null) return { line, rule: { kind: "ip", value: hostname } };
const rule: CompanyScopeRule = { kind: "domain", value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
if (/icp|备案/i.test(value)) return { line, rule: { kind: "icp", value } };
return { line, rule: { kind: "keyword", value } };
}
export function parseCompanyScopeText(
value: string,
options: { preservedRules?: CompanyScopeRule[]; maxRules?: number } = {},
): ParsedCompanyScopeText {
const nonEmpty = value
.split(/\r?\n/)
.map((raw, index) => ({ raw, line: index + 1 }))
.filter(({ raw }) => raw.trim());
const preserved = new Map<string, CompanyScopeRule[]>();
for (const rule of options.preservedRules ?? []) {
const key = rule.value.trim();
preserved.set(key, [...(preserved.get(key) ?? []), rule]);
}
const issues = nonEmpty.map(({ raw, line }) => {
const key = raw.trim();
const queue = preserved.get(key);
const preservedRule = queue?.shift();
return classifyCompanyScopeLine(raw, line, preservedRule);
});
const maxRules = options.maxRules ?? MAX_COMPANY_SCOPE_RULES;
if (nonEmpty.length > maxRules) {
issues.push({
line: nonEmpty[Math.max(0, maxRules)]?.line ?? 1,
error: `最多允许 ${maxRules} 条范围,当前 ${nonEmpty.length} 条`,
});
}
return {
rules: issues.flatMap((item) => (item.rule ? [item.rule] : [])),
errors: issues.flatMap((item) => (item.error ? [{ line: item.line, error: item.error }] : [])),
};
}
export function normalizeCompanyScopeValue(rule: CompanyScopeRule): string {
const value = rule.value.trim();
if (rule.kind === "domain") return domainHostname(value) ?? value.toLowerCase();
if (rule.kind === "icp") return value.toLowerCase().replace(/\s/gu, "");
if (rule.kind === "keyword") return value.toLowerCase().split(/\s+/u).filter(Boolean).join(" ");
return value.toLowerCase();
}
+49 -3
View File
@@ -7,8 +7,8 @@ import type {
Activity, Agent, AgentDetail, Asset, Audit, Company, Conversation,
ConvTokenSummary, DailyTokenBucket, Edge, Finding, InterceptApprovalRow,
InterceptPending, InterceptRule, LLMPoolStatus, LLMProfile, MCPServer, MCPTool, PromptVar,
PromptVersion, Settings, SkillCall, SkillItem, MissingSkill, Stats, TaskNode, Task, TokenTotal,
TokenUsage, Tool, TrafficResp, TrafficDetail, TrafficHost, LLMTask,
PromptVersion, Settings, SessionTokenUsage, SkillCall, SkillItem, MissingSkill, Stats, TaskNode, Task, TokenTotal,
TaskTemplate, TokenUsage, Tool, TrafficResp, TrafficDetail, TrafficHost, LLMTask,
} from "@/lib/types";
const T = (iso: string) => iso; // readability helper for timestamps
@@ -93,6 +93,25 @@ export const tasks: Task[] = [
},
];
export const taskTemplates: TaskTemplate[] = [
{
id: 1,
name: "外部 Web 渗透",
description: "对目标互联网暴露面开展黑盒渗透测试,覆盖站点、接口和常见管理入口。",
goal: "识别并验证可造成未授权访问、敏感数据泄露或服务器失陷的安全问题。",
created_at: T("2026-07-20T08:00:00Z"),
updated_at: T("2026-07-25T08:00:00Z"),
},
{
id: 2,
name: "API 越权专项",
description: "围绕目标 API 的身份认证、对象级授权和角色边界开展专项测试。",
goal: "确认是否存在 IDOR、水平越权、垂直越权及批量数据访问风险。",
created_at: T("2026-07-19T08:00:00Z"),
updated_at: T("2026-07-24T08:00:00Z"),
},
];
export const getTask = (id: string) => tasks.find((t) => t.id === id);
const grandTotal: TokenTotal = tasks.reduce(
@@ -516,6 +535,33 @@ export const tokenWorkers: TokenUsage[] = [
{ worker: "work#13", input_tokens: 144000, output_tokens: 9200, cache_read_tokens: 96000, cache_write_tokens: 14000 },
];
// Mirrors TokenStatsBySession: completed usage is keyed by the stable UI session
// (main | plan | intent:<id>) instead of the reusable work#N executor name.
export const tokenSessions: SessionTokenUsage[] = (() => {
const totals = new Map<string, SessionTokenUsage>();
for (const item of activity) {
if (item.kind !== "result") continue;
let session = "";
if (item.worker === "mainagent") session = "main";
else if (item.worker === "planner") session = "plan";
else if (item.intent_id) session = `intent:${item.intent_id}`;
if (!session) continue;
const current = totals.get(session) ?? {
session,
input_tokens: 0,
output_tokens: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
};
current.input_tokens += item.input_tokens ?? 0;
current.output_tokens += item.output_tokens ?? 0;
current.cache_read_tokens += item.cache_read_tokens ?? 0;
current.cache_write_tokens += item.cache_write_tokens ?? 0;
totals.set(session, current);
}
return [...totals.values()];
})();
export const tokenTotal: TokenTotal = grandTotal;
export const dailyTokens: DailyTokenBucket[] = Array.from({ length: 30 }, (_, i) => {
@@ -755,7 +801,7 @@ export const interceptHistory: InterceptApprovalRow[] = [
// ── Conversations (chat) ─────────────────────────────────────────────────────
export const conversations: Conversation[] = [
{ id: 1, agent_key: "mainagent", title: "acme 后台切入点讨论", llm_profile_id: 1, created_at: T("2026-07-25T10:00:00Z"), updated_at: T("2026-07-26T03:40:00Z") },
{ id: 1, agent_key: "mainagent", title: "acme 后台切入点讨论", llm_profile_id: 1, pinned: true, pinned_at: T("2026-07-26T04:00:00Z"), created_at: T("2026-07-25T10:00:00Z"), updated_at: T("2026-07-26T03:40:00Z") },
{ id: 2, agent_key: "mainagent", title: "API 越权面梳理", llm_profile_id: 2, created_at: T("2026-07-24T16:00:00Z"), updated_at: T("2026-07-24T17:20:00Z") },
{ id: 3, agent_key: "mainagent", title: "内网横向 · DMZ→域控攻击路径规划", llm_profile_id: 1, created_at: T("2026-07-25T15:45:00Z"), updated_at: T("2026-07-26T00:25:00Z") },
];
+473 -22
View File
@@ -2,7 +2,24 @@
// 未命中的一律返回安全默认([] / {} / {ok:true}),保证任何页面都不崩。
// 只在 NEXT_PUBLIC_MOCK=1 时经由 api.ts 的 http() 短路进入这里。
import type { Task } from "../types";
import {
classifyCompanyScopeLine,
companyScopeRuleError,
isCompanyScopeKind,
MAX_COMPANY_SCOPE_RULES,
normalizeCompanyScopeValue,
} from "../company-scope";
import type {
Asset,
BatchControlItem,
Company,
CompanyScopeRule,
Conversation,
ScopeRow,
Task,
TaskLLMResolution,
TaskTemplate,
} from "../types";
import * as D from "./data";
const delay = (ms = 120) => new Promise((r) => setTimeout(r, ms));
@@ -12,8 +29,181 @@ const delay = (ms = 120) => new Promise((r) => setTimeout(r, ms));
const mockTasks = structuredClone(D.tasks);
const mockFindings = structuredClone(D.findings);
const mockLLMRecords = structuredClone(D.llmRecords);
const mockTaskTemplates = structuredClone(D.taskTemplates);
const mockConversations = structuredClone(D.conversations);
const mockIntents = structuredClone(D.intents);
const mockCompanies = structuredClone(D.companies);
const mockAssets = structuredClone(D.assets);
let mockActiveTask = D.ACTIVE_TASK;
function mockAssetCounts(): Record<string, number> {
return mockAssets.reduce<Record<string, number>>((counts, asset) => {
counts[asset.type] = (counts[asset.type] ?? 0) + 1;
return counts;
}, {});
}
function mockScopeRows(
companyID: number,
input: unknown,
existing: ScopeRow[] = [],
): { rows: ScopeRow[]; invalid: number; skipped: number } {
if (!Array.isArray(input)) return { rows: [], invalid: 0, skipped: 0 };
if (input.length > MAX_COMPANY_SCOPE_RULES) throw new Error(`企业范围最多支持 ${MAX_COMPANY_SCOPE_RULES} 条`);
let nextID =
mockCompanies.flatMap((company) => company.scope ?? []).reduce((max, row) => Math.max(max, row.id), 0) + 1;
const rows: ScopeRow[] = [];
let invalid = 0;
let skipped = 0;
const keys = new Set(
existing.map((row) => `${row.kind}|${row.domain ?? row.net ?? row.value ?? row.raw.trim().toLowerCase()}`),
);
for (const [index, candidate] of input.entries()) {
let rule: CompanyScopeRule | undefined;
if (typeof candidate === "string") {
rule = classifyCompanyScopeLine(candidate, index + 1).rule;
} else if (candidate && typeof candidate === "object") {
const item = candidate as { kind?: unknown; value?: unknown };
const value = String(item.value ?? "").trim();
if (item.kind === undefined || item.kind === "") rule = classifyCompanyScopeLine(value, index + 1).rule;
else if (isCompanyScopeKind(item.kind)) rule = { kind: item.kind, value };
}
if (!rule || companyScopeRuleError(rule)) {
invalid++;
continue;
}
const normalized = normalizeCompanyScopeValue(rule);
const row: ScopeRow = { id: nextID++, company_id: companyID, kind: rule.kind, raw: rule.value.trim() };
if (rule.kind === "domain") row.domain = normalized;
else if (rule.kind === "ip") row.net = `${normalized}/${normalized.includes(":") ? 128 : 32}`;
else if (rule.kind === "cidr") row.net = normalized;
else row.value = normalized;
const key = `${row.kind}|${row.domain ?? row.net ?? row.value}`;
if (keys.has(key)) {
skipped++;
continue;
}
keys.add(key);
rows.push(row);
}
if (existing.length + rows.length > MAX_COMPANY_SCOPE_RULES) {
throw new Error(`企业范围规则过多:每个企业最多 ${MAX_COMPANY_SCOPE_RULES} 条`);
}
return { rows, invalid, skipped };
}
function mockProfileResolution(profileID: number | undefined, source: TaskLLMResolution["source"]): TaskLLMResolution {
const profile = D.llmProfiles.find((item) => Number(item.id) === profileID);
if (!profile) {
return { name: "", format: "", model: "", source, available: false, reason: "LLM 配置不存在" };
}
return {
profile_id: Number(profile.id),
name: profile.name,
format: profile.format,
model: profile.model,
source,
available: Boolean(profile.api_key_hint),
reason: profile.api_key_hint ? undefined : "LLM 配置未设置 API Key",
};
}
// Mirrors the backend precedence in server/task_resolution.go:
// Agent 绑定 → 任务 LLM 配置链 → 全局配置 → 环境配置。
function mockRoleResolution(task: Task, agentKey: "mainagent" | "planner" | "worker"): TaskLLMResolution {
const agent = D.agents.find((item) => item.key === agentKey);
if (agent?.llm_profile_id) {
const bound = mockProfileResolution(agent.llm_profile_id, "agent_binding");
if (bound.available) return bound;
}
if (task.llm_profile_ids?.length) {
if (task.active_llm_profile_id === undefined) {
return {
name: "",
format: "",
model: "",
source: "task_chain",
available: false,
reason: "任务 LLM 配置链额度已耗尽",
};
}
return mockProfileResolution(task.active_llm_profile_id, "task_chain");
}
const globalProfile = D.llmProfiles.find((item) => item.is_default);
if (globalProfile) return mockProfileResolution(Number(globalProfile.id), "global_profile");
return {
name: "全局配置",
format: D.llmConfig.provider,
model: D.llmConfig.model,
source: "environment",
available: true,
};
}
function bodyIDs(value: unknown): string[] {
if (!Array.isArray(value)) return [];
const ids = [...new Set(value.map((item) => String(item)).filter(Boolean))];
if (ids.length > 100) throw new Error("批量操作最多支持 100 个 ID");
return ids;
}
type MockIntentControlResult = { ok: boolean; state?: string; error?: string };
function controlMockIntent(id: string, action: "pause" | "resume"): MockIntentControlResult {
const intent = mockIntents.find((item) => item.id === id);
if (!intent) return { ok: false, error: "意图不存在" };
const requiredState = action === "pause" ? "running" : "paused";
if (intent.inherited || intent.state !== requiredState) {
return { ok: false, state: intent.state, error: "Worker 状态已变化" };
}
intent.state = action === "pause" ? "paused" : "open";
return { ok: true, state: intent.state };
}
function controlMockTask(id: string, action: "pause" | "resume"): BatchControlItem {
const task = mockTasks.find((item) => item.id === id);
if (!task) return { id, ok: false, error: "任务不存在" };
if (task.status === "done" || task.status === "failed" || task.status === "timeout") {
return { id, ok: false, status: task.status, error: "终态任务不可控制" };
}
if (action === "pause") {
if (task.paused || task.status === "paused") {
return { id, ok: false, status: task.status, error: "任务已经暂停" };
}
task.paused = true;
task.queued = false;
task.status = "paused";
task.engine_mode = "paused";
} else {
if (!task.paused && task.status !== "paused") {
return { id, ok: false, status: task.status, error: "任务未暂停" };
}
task.paused = false;
task.queued = false;
task.status = "running";
task.engine_mode = "exploring";
}
return { id, ok: true, status: task.status, queued: false };
}
function sortMockConversations() {
mockConversations.sort((a, b) => {
const aPinned = a.pinned_at ? 1 : 0;
const bPinned = b.pinned_at ? 1 : 0;
if (aPinned !== bPinned) return bPinned - aPinned;
const aTime = a.pinned_at ?? a.updated_at;
const bTime = b.pinned_at ?? b.updated_at;
return bTime.localeCompare(aTime) || b.id - a.id;
});
}
function normalizedTemplateName(value: unknown): string {
return String(value ?? "")
.trim()
.split(/\s+/)
.join(" ");
}
function parseBody(body?: BodyInit | null): Record<string, unknown> {
if (typeof body !== "string") return {};
try {
@@ -50,6 +240,7 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
const now = new Date();
const profileIDs = [...((b.llm_profile_ids as number[] | undefined) ?? [])];
const sourceTaskIDs = [...((b.source_task_ids as string[] | undefined) ?? [])];
const companyIDs = [...((b.company_ids as number[] | undefined) ?? [])];
const created: Task = {
id,
description: String(b.description ?? "新任务"),
@@ -70,12 +261,64 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
active_llm_profile_id: profileIDs[0],
llm_failover_state: profileIDs.length ? "ready" : "default",
source_task_ids: sourceTaskIDs,
company_ids: companyIDs,
};
for (const item of mockTasks) item.active = false;
mockTasks.unshift(created);
mockActiveTask = id;
return created;
}
if (path === "/task-templates" && m === "GET") return { templates: mockTaskTemplates };
if (path === "/task-templates" && m === "POST") {
const now = new Date().toISOString();
const name = normalizedTemplateName(b.name);
const description = String(b.description ?? "").trim();
const goal = String(b.goal ?? "").trim();
if (!name || !description || !goal) throw new Error("请填写模板名称、描述和目标");
if (
mockTaskTemplates.some((template) => normalizedTemplateName(template.name).toLowerCase() === name.toLowerCase())
) {
throw new Error("模板名称已存在");
}
const nextID = mockTaskTemplates.reduce((max, template) => Math.max(max, template.id), 0) + 1;
const created: TaskTemplate = {
id: nextID,
name,
description,
goal,
created_at: now,
updated_at: now,
};
mockTaskTemplates.unshift(created);
return created;
}
if (seg[0] === "task-templates" && seg.length === 2 && m === "PATCH") {
const template = mockTaskTemplates.find((item) => item.id === Number(seg[1]));
if (!template) return {};
const name = typeof b.name === "string" ? normalizedTemplateName(b.name) : template.name;
const description = typeof b.description === "string" ? b.description.trim() : template.description;
const goal = typeof b.goal === "string" ? b.goal.trim() : template.goal;
if (!name || !description || !goal) throw new Error("请填写模板名称、描述和目标");
if (
mockTaskTemplates.some(
(item) => item.id !== template.id && normalizedTemplateName(item.name).toLowerCase() === name.toLowerCase(),
)
) {
throw new Error("模板名称已存在");
}
template.name = name;
template.description = description;
template.goal = goal;
template.updated_at = new Date().toISOString();
mockTaskTemplates.sort((a, b) => b.updated_at.localeCompare(a.updated_at) || b.id - a.id);
return template;
}
if (seg[0] === "task-templates" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockTaskTemplates.findIndex((item) => item.id === id);
if (index >= 0) mockTaskTemplates.splice(index, 1);
return { deleted: id };
}
if (seg[0] === "tasks" && seg.length === 2 && m === "DELETE") {
const id = seg[1];
const index = mockTasks.findIndex((item) => item.id === id);
@@ -88,6 +331,14 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
mockFindings.splice(i, 1);
findingsDeleted++;
}
} else {
// PostgreSQL uses ON DELETE SET NULL for retained findings. Keep the mock
// grouped view consistent by moving them into the unassigned/deleted bucket.
for (const finding of mockFindings) {
if (finding.task_id !== id) continue;
finding.task_id = undefined;
finding.task_description = "";
}
}
let llmRecordsDeleted = 0;
@@ -133,7 +384,52 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
reopened_intents: 0,
};
}
if (seg[0] === "tasks" && seg[2] === "control") return { id: seg[1], paused: b.action === "pause" };
if (seg[0] === "tasks" && seg[2] === "llm" && seg[3] === "resolution" && m === "GET") {
const task = mockTasks.find((item) => item.id === seg[1]);
if (!task) throw new Error("任务不存在");
return {
mainagent: mockRoleResolution(task, "mainagent"),
planner: mockRoleResolution(task, "planner"),
worker: mockRoleResolution(task, "worker"),
};
}
if (path === "/tasks/control/batch" && m === "POST") {
const action = b.action === "resume" ? "resume" : "pause";
return { items: bodyIDs(b.task_ids).map((id) => controlMockTask(id, action)) };
}
if (seg[0] === "tasks" && seg[2] === "intents" && seg[4] === "control" && m === "POST") {
const id = seg[3];
if (b.action === "cancel") {
const index = mockIntents.findIndex((item) => item.id === id);
const intent = mockIntents[index];
if (!intent || intent.inherited || (intent.state !== "running" && intent.state !== "paused")) {
throw new Error("Worker 状态已变化");
}
mockIntents.splice(index, 1);
return {
id: Number(id.replace(/\D/g, "")) || 0,
state: "cancelled",
deleted: {
intents: 1,
facts: 1,
findings: 1,
activities: D.activity.filter((item) => item.intent_id === id).length,
},
};
}
const action = b.action === "resume" ? "resume" : "pause";
const result = controlMockIntent(id, action);
if (!result.ok) throw new Error(result.error ?? "Worker 状态已变化");
return { id: Number(id.replace(/\D/g, "")) || 0, state: result.state };
}
if (seg[0] === "tasks" && seg.length === 3 && seg[2] === "control" && m === "POST") {
const action = b.action === "resume" ? "resume" : "pause";
const result = controlMockTask(seg[1], action);
if (!result.ok) throw new Error(result.error ?? "任务状态已变化");
const task = mockTasks.find((item) => item.id === seg[1]);
return { id: seg[1], paused: Boolean(task?.paused), queued: Boolean(task?.queued), status: task?.status ?? "" };
}
if (seg[0] === "tasks" && seg[2] === "chat" && seg[3] === "status") return { running: false };
if (seg[0] === "tasks" && seg[2] === "chat" && seg[3] === "stop") return { status: "stopped" };
if (path === "/active") {
const id = String(b.id ?? mockActiveTask);
@@ -227,21 +523,72 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
}
// ── assets ──
if (path === "/assets/counts") return D.assetCounts;
if (path === "/assets/counts") return mockAssetCounts();
if (path === "/assets" && m === "GET") {
const type = q.get("type") ?? "";
const list = type ? D.assets.filter((a) => a.type === type) : D.assets;
const list = type ? mockAssets.filter((a) => a.type === type) : mockAssets;
const limit = Number(q.get("limit") ?? 50);
const offset = Number(q.get("offset") ?? 0);
return { count: list.length, total: list.length, assets: list.slice(offset, offset + limit) };
}
if (path === "/assets" && m === "DELETE") return { deleted: (b.ids as unknown[])?.length ?? 0 };
if (path === "/assets" && m === "DELETE") {
const ids = new Set(Array.isArray(b.ids) ? b.ids.map(Number) : []);
let deleted = 0;
for (let index = mockAssets.length - 1; index >= 0; index--) {
if (!ids.has(mockAssets[index].id)) continue;
mockAssets.splice(index, 1);
deleted++;
}
return { deleted };
}
// ── companies ──
if (path === "/companies" && m === "GET") return D.companies;
if (path === "/companies" && m === "POST") return { id: 2, created: true, scope_added: 0 };
if (seg[0] === "companies" && seg[2] === "scope") return { added: 0, skipped: 0, invalid: 0 };
if (seg[0] === "companies" && seg.length === 2 && m === "DELETE") return { deleted: 1, assets_deleted: 0 };
if (path === "/companies" && m === "GET") return structuredClone(mockCompanies);
if (path === "/companies" && m === "POST") {
const name = String(b.name ?? "").trim();
if (!name) throw new Error("企业名称不能为空");
if (mockCompanies.some((company) => company.name.toLowerCase() === name.toLowerCase())) {
throw new Error("企业已存在");
}
const id = mockCompanies.reduce((max, company) => Math.max(max, company.id), 0) + 1;
const scopeResult = mockScopeRows(id, b.scope);
const company: Company = { id, name, asset_count: 0, scope: scopeResult.rows };
mockCompanies.push(company);
return {
id,
created: true,
scope_added: scopeResult.rows.length,
scope_skipped: scopeResult.skipped,
scope_invalid: scopeResult.invalid,
};
}
if (seg[0] === "companies" && seg[2] === "scope" && m === "POST") {
const company = mockCompanies.find((item) => item.id === Number(seg[1]));
if (!company) throw new Error("企业不存在");
const reset = b.reset === true;
const scopeResult = mockScopeRows(company.id, b.scope, reset ? [] : (company.scope ?? []));
if (reset && scopeResult.invalid > 0) throw new Error("企业范围包含无效规则,未覆盖原有范围");
company.scope = reset ? scopeResult.rows : [...(company.scope ?? []), ...scopeResult.rows];
return { added: scopeResult.rows.length, skipped: scopeResult.skipped, invalid: scopeResult.invalid };
}
if (seg[0] === "companies" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockCompanies.findIndex((item) => item.id === id);
if (index < 0) throw new Error("企业不存在");
mockCompanies.splice(index, 1);
let assetsDeleted = 0;
for (let assetIndex = mockAssets.length - 1; assetIndex >= 0; assetIndex--) {
const asset: Asset = mockAssets[assetIndex];
if (asset.company_id !== id) continue;
if (b.delete_assets === true) {
mockAssets.splice(assetIndex, 1);
assetsDeleted++;
} else {
delete asset.company_id;
}
}
return { deleted: 1, assets_deleted: assetsDeleted };
}
// ── exploration ──
if (path === "/exploration/frontier") return D.frontier;
@@ -267,6 +614,72 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
tasks,
};
}
if (path === "/exploration/findings/groups") {
const severityOrder = { critical: 4, high: 3, medium: 2, low: 1 } as const;
let list = mockFindings.slice();
const filterSeverity = q.get("severity");
const filterStatus = q.get("status");
const filterVulnclass = q.get("vulnclass");
const filterTask = q.get("task_id");
if (filterSeverity) list = list.filter((finding) => finding.severity === filterSeverity);
if (filterStatus) list = list.filter((finding) => finding.status === filterStatus);
if (filterVulnclass) list = list.filter((finding) => finding.vulnclass === filterVulnclass);
if (filterTask === "__unassigned__") list = list.filter((finding) => !finding.task_id);
else if (filterTask) list = list.filter((finding) => finding.task_id === filterTask);
const grouped = new Map<string, typeof list>();
for (const finding of list) {
const key = finding.task_id ?? "__unassigned__";
grouped.set(key, [...(grouped.get(key) ?? []), finding]);
}
const groups = Array.from(grouped, ([key, items]) => {
const owner = mockTasks.find((candidate) => candidate.id === key);
return {
task_id: key === "__unassigned__" ? null : key,
task_description: owner?.description ?? items[0]?.task_description ?? "",
task_status: owner?.status ?? "",
count: items.length,
critical: items.filter((finding) => finding.severity === "critical").length,
high: items.filter((finding) => finding.severity === "high").length,
medium: items.filter((finding) => finding.severity === "medium").length,
low: items.filter((finding) => finding.severity === "low").length,
last_found_at: items.reduce((latest, finding) => (finding.ts > latest ? finding.ts : latest), ""),
max_severity: Math.max(...items.map((finding) => severityOrder[finding.severity])),
};
});
groups.sort((left, right) =>
q.get("sort") === "severity"
? right.max_severity - left.max_severity || right.last_found_at.localeCompare(left.last_found_at)
: right.last_found_at.localeCompare(left.last_found_at),
);
const rawPage = Number(q.get("page") ?? 1);
const rawPageSize = Number(q.get("limit") ?? 10);
const page = Number.isFinite(rawPage) && rawPage > 0 ? Math.floor(rawPage) : 1;
const pageSize = Number.isFinite(rawPageSize) && rawPageSize > 0 ? Math.min(100, Math.floor(rawPageSize)) : 10;
return {
items: groups.slice((page - 1) * pageSize, page * pageSize),
total: groups.length,
finding_total: list.length,
page,
page_size: pageSize,
};
}
if (seg[0] === "exploration" && seg[1] === "findings" && seg[3] === "deepen" && m === "POST") {
const finding = mockFindings.find((candidate) => candidate.id === seg[2]);
const description = String(b.description ?? "").trim();
if (!description) throw new Error("description is required");
if ([...description].length > 4000) throw new Error("description must be at most 4000 characters");
if (!finding) throw new Error("finding not found");
if (!finding.task_id || !mockTasks.some((candidate) => candidate.id === finding.task_id)) {
throw new Error("finding origin task or node is no longer available");
}
return {
task_id: finding.task_id,
intent_id: `mock-deepen-${Date.now()}`,
state: "open",
queued: false,
};
}
// 单条 finding:GET 详情 / PATCH 改状态/严重度/名称/类别(demo 直接改内存对象)。
if (seg[0] === "exploration" && seg[1] === "findings" && seg.length === 3 && seg[2] !== "stats") {
const f = mockFindings.find((x) => x.id === seg[2]);
@@ -320,14 +733,17 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (fSev) list = list.filter((f) => f.severity === fSev);
if (fStatus) list = list.filter((f) => f.status === fStatus);
if (fVuln) list = list.filter((f) => f.vulnclass === fVuln);
if (fTask) list = list.filter((f) => f.task_id === fTask);
if (fTask === "__unassigned__") list = list.filter((f) => !f.task_id);
else if (fTask) list = list.filter((f) => f.task_id === fTask);
list.sort((a, b) =>
q.get("sort") === "severity"
? sev[b.severity] - sev[a.severity] || +new Date(b.ts) - +new Date(a.ts)
: +new Date(b.ts) - +new Date(a.ts),
);
const page = Number(q.get("page") ?? 1);
const pageSize = Number(q.get("limit") ?? 20);
const rawPage = Number(q.get("page") ?? 1);
const rawPageSize = Number(q.get("limit") ?? 20);
const page = Number.isFinite(rawPage) && rawPage > 0 ? Math.floor(rawPage) : 1;
const pageSize = Number.isFinite(rawPageSize) && rawPageSize > 0 ? Math.min(200, Math.floor(rawPageSize)) : 20;
return {
items: list.slice((page - 1) * pageSize, page * pageSize).map(withFid),
total: list.length,
@@ -339,13 +755,20 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (q.has("page")) {
const before = Number(q.get("before") ?? 0);
const limit = Math.max(1, Number(q.get("limit") ?? 300));
let list = D.intents;
let list = mockIntents;
if (before > 0) list = list.filter((intent) => Number(intent.id.replace(/\D/g, "") || intent.id) < before);
return { items: list.slice(0, limit), has_more: list.length > limit };
}
return D.intents;
return mockIntents;
}
if (path === "/exploration/tokens") {
const selectedTask = task ? mockTasks.find((item) => item.id === task) : undefined;
return {
workers: D.tokenWorkers,
sessions: D.tokenSessions,
total: selectedTask?.tokens ?? D.tokenTotal,
};
}
if (path === "/exploration/tokens") return { workers: D.tokenWorkers, total: D.tokenTotal };
if (path === "/exploration/graph") return D.explorationGraph;
if (path === "/exploration/activity" && seg.length === 2) {
const since = Number(q.get("since") ?? 0);
@@ -425,15 +848,43 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re
if (seg[0] === "agents" && seg[2] === "visibility" && m === "GET") return D.agentDetail(seg[1]).visibility;
// ── conversations ──
if (path === "/conversations" && m === "GET") return { conversations: D.conversations };
if (path === "/conversations" && m === "POST")
return {
id: 3,
if (path === "/conversations" && m === "GET") {
sortMockConversations();
return { conversations: structuredClone(mockConversations) };
}
if (path === "/conversations" && m === "POST") {
const now = new Date().toISOString();
const title = String(b.title ?? "").trim() || "新对话";
const conversation: Conversation = {
id: mockConversations.reduce((max, item) => Math.max(max, item.id), 0) + 1,
agent_key: String(b.agent_key ?? "mainagent"),
title: String(b.title ?? "新会话"),
created_at: "2026-07-26T04:00:00Z",
updated_at: "2026-07-26T04:00:00Z",
title,
llm_profile_id: typeof b.llm_profile_id === "number" ? b.llm_profile_id : undefined,
pinned: false,
created_at: now,
updated_at: now,
};
mockConversations.unshift(conversation);
return structuredClone(conversation);
}
if (seg[0] === "conversations" && seg.length === 2 && m === "PATCH") {
const conversation = mockConversations.find((item) => item.id === Number(seg[1]));
if (!conversation) return {};
if (typeof b.title === "string") conversation.title = b.title.trim();
if (typeof b.pinned === "boolean") {
conversation.pinned = b.pinned;
conversation.pinned_at = b.pinned ? (conversation.pinned_at ?? new Date().toISOString()) : null;
}
conversation.updated_at = new Date().toISOString();
sortMockConversations();
return structuredClone(conversation);
}
if (seg[0] === "conversations" && seg.length === 2 && m === "DELETE") {
const id = Number(seg[1]);
const index = mockConversations.findIndex((item) => item.id === id);
if (index >= 0) mockConversations.splice(index, 1);
return { deleted: id };
}
if (seg[0] === "conversations" && seg[2] === "messages" && seg.length === 3 && m === "GET") {
const items = D.conversationMessages[Number(seg[1])] ?? [];
return { items, cursor: items.length ? items[items.length - 1].seq : 0, running: false };
+84 -4
View File
@@ -30,6 +30,16 @@ export interface Task {
llm_failover_state?: "default" | "ready" | "chain_exhausted" | string;
llm_failover_reason?: string;
source_task_ids?: string[]; // directly related tasks inherited as read-only context
company_ids?: number[]; // associated company asset scopes available as task context
}
export interface TaskTemplate {
id: number;
name: string;
description: string;
goal: string;
created_at: string;
updated_at: string;
}
export interface DeleteTaskOptions {
@@ -237,13 +247,22 @@ export interface TaskScopeRow {
reason?: string;
}
export type CompanyScopeKind = "domain" | "ip" | "cidr" | "icp" | "keyword";
// 新增企业时提交的结构化资产范围规则。
export interface CompanyScopeRule {
kind: CompanyScopeKind;
value: string;
}
// 公司资产范围规则的一条(归属唯一真值来源)。
export interface ScopeRow {
id: number;
company_id: number;
kind: "domain" | "ip" | "cidr";
kind: CompanyScopeKind;
domain?: string; // kind=domain 时有值
net?: string; // kind=ip|cidr 时有值
value?: string; // kind=icp|keyword 时可能由后端直接返回
raw: string; // 原始用户输入,用于显示和回填
reason?: string;
}
@@ -254,13 +273,13 @@ export interface Company {
name: string;
logo?: string; // 远程图标 URL;空则前端用名称首字母
asset_count: number;
scope: ScopeRow[];
scope?: ScopeRow[];
}
// ---- Exploration graph (per task) ----
export type ExploreKind = "task" | "begin" | "goal" | "intent" | "fact" | "finding" | "hint";
export type GoalState = "open" | "met" | "abandoned";
export type IntentState = "open" | "running" | "done" | "blocked" | "exhausted" | "stopped";
export type IntentState = "open" | "running" | "paused" | "done" | "blocked" | "exhausted" | "stopped";
export type FindingState = "confirmed" | "dismissed";
export type HintState = "active" | "consumed";
export type ExploreRel = "spawns" | "derived_from" | "yields" | "proves";
@@ -326,6 +345,33 @@ export interface FindingsPage {
page_size: number;
}
export interface FindingGroup {
task_id: string | number | null;
task_description: string;
task_status: string;
count: number;
critical: number;
high: number;
medium: number;
low: number;
last_found_at: string;
}
export interface FindingGroupsPage {
items: FindingGroup[];
total: number;
finding_total: number;
page: number;
page_size: number;
}
export interface FindingDeepenResponse {
task_id: string;
intent_id: string;
state: IntentState;
queued: boolean;
}
// FindingStats 是发现全表聚合(统计卡 + 漏洞类型下拉),服务端计算,不受分页影响。
export interface FindingStats {
total: number;
@@ -341,7 +387,7 @@ export interface FindingStats {
// FindingTaskOption 是发现页「按任务」筛选下拉的一项:有漏洞的任务(描述为空表示任务已删除,
// 前端回退展示 id)及其漏洞条数。
export interface FindingTaskOption {
id: number;
id: string | number;
description: string;
count: number;
}
@@ -417,6 +463,22 @@ export interface LLMTransition {
next?: LLMAuditProfile;
}
export interface TaskLLMResolution {
profile_id?: number;
name: string;
format: string;
model: string;
source: "task_chain" | "agent_binding" | "global_profile" | "environment" | "global";
available: boolean;
reason?: string;
}
export interface TaskLLMResolutions {
mainagent: TaskLLMResolution;
planner: TaskLLMResolution;
worker: TaskLLMResolution;
}
// ---- Agent triggers (P3 调度,仅自定义 agent) ----
export interface AgentTrigger {
id: number;
@@ -444,6 +506,8 @@ export interface Conversation {
agent_key: string;
title: string;
llm_profile_id?: number;
pinned?: boolean;
pinned_at?: string | null;
created_at: string;
updated_at: string;
}
@@ -479,6 +543,22 @@ export interface TokenUsage {
cache_write_tokens: number;
}
export interface SessionTokenUsage {
session: string;
input_tokens: number;
output_tokens: number;
cache_read_tokens: number;
cache_write_tokens: number;
}
export interface BatchControlItem {
id: string;
ok: boolean;
status?: string;
queued?: boolean;
error?: string;
}
// Whole-task (all agents) token aggregate.
export interface TokenTotal {
input_tokens: number;