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:
commit
f0c2b43802
78 files changed
+12765
-2651
No files matched your search
@@ -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
@@ -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
|
||||
@@ -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
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 更新升级
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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))
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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 导出(导出当前筛选/全部)。结果按严重等级降序、
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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{
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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():
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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 = ""
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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, ©)
|
||||
}
|
||||
}
|
||||
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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
@@ -122,7 +122,7 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 500; i++ {
|
||||
id := int64(offset*1000 + i + 1)
|
||||
task.setLLMState(&id, &id, []int64{id, id + 1}, "ready", "")
|
||||
task.setLLMState(&id, &id, []int64{id, id + 1}, int64(i), "ready", "")
|
||||
}
|
||||
}(writer)
|
||||
}
|
||||
@@ -146,6 +146,24 @@ func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMStateRejectsOlderRevision(t *testing.T) {
|
||||
t.Parallel()
|
||||
task := &Task{ID: "7"}
|
||||
current := int64(22)
|
||||
stale := int64(11)
|
||||
if !task.setLLMState(¤t, ¤t, []int64{current}, 2, "ready", "") {
|
||||
t.Fatal("initial LLM state was rejected")
|
||||
}
|
||||
if task.setLLMState(&stale, &stale, []int64{stale}, 1, "ready", "stale") {
|
||||
t.Fatal("older LLM snapshot was accepted")
|
||||
}
|
||||
state := task.llmStateSnapshot()
|
||||
if state.ChainRevision != 2 || state.ActiveID == nil || *state.ActiveID != current ||
|
||||
len(state.ProfileIDs) != 1 || state.ProfileIDs[0] != current || state.FailoverReason != "" {
|
||||
t.Fatalf("older LLM snapshot overwrote current state: %+v", state)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) {
|
||||
t.Parallel()
|
||||
quota := errors.New("openai: status 402: insufficient_quota")
|
||||
@@ -381,7 +399,7 @@ func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) {
|
||||
engine.started.Store("7", true)
|
||||
task := &Task{ID: "7", Status: "running"}
|
||||
profileID := int64(11)
|
||||
task.setLLMState(&profileID, nil, []int64{profileID}, "chain_exhausted", "quota exhausted")
|
||||
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
|
||||
s := &Server{engine: engine}
|
||||
if got := s.resolvedTaskStatus(task); got != "running" {
|
||||
t.Fatalf("exhausted started status=%q, want running", got)
|
||||
@@ -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{}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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(""));
|
||||
}, []);
|
||||
|
||||
|
||||
@@ -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();
|
||||
}}
|
||||
|
||||
@@ -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
@@ -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
@@ -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>{" "}
|
||||
的所有流量记录(含请求/响应原文),此操作不可撤销。
|
||||
</>
|
||||
)}
|
||||
|
||||
@@ -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}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -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
@@ -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}`),
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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;
|
||||
|
||||
Reference in new issue
Block a user