chore: ARTEX 初始化
ARTEX — AI 自主渗透测试系统（Go 后端 + Next.js 前端，单二进制内嵌）。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Autumn-27andClaude Opus 4.8 committed
commit
0ef8561411
287 files changed
+70178
No files matched your search
@@ -0,0 +1,689 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/enrich"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// Task is one engagement: a description + goal + its own exploration store,
|
||||
// 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"`
|
||||
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
|
||||
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态；空/其它则由运行态推导)
|
||||
// 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
FirstRunAt int64 `json:"first_run_at,omitempty"`
|
||||
DeadlineAt int64 `json:"deadline_at,omitempty"`
|
||||
Store *pgdb.ExplorationStore `json:"-"`
|
||||
Guard *guard.Guard `json:"-"`
|
||||
notify chan struct{}
|
||||
|
||||
// doneIntents accumulates the ids of intents that completed since the last
|
||||
// planning round consumed them. The debounce coalesces a burst of completions
|
||||
// into one round, so several ids may pile up before drainDone() clears them.
|
||||
doneMu sync.Mutex
|
||||
doneIntents []int64
|
||||
}
|
||||
|
||||
// Manager owns the PostgreSQL data source (asset graph + every task's exploration
|
||||
// graph + config) and the in-memory set of task handles.
|
||||
type Manager struct {
|
||||
dir string
|
||||
pg *pgdb.DB
|
||||
assets *pgdb.AssetStore
|
||||
traffic *traffic.Traffic // process-wide recording proxy (may be nil)
|
||||
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）
|
||||
// 联网搜索开关与来源（默认关；settings.web_search_*）。brave-free 需要 braveKey；tavily 需要 tavilyKey。
|
||||
// webSearchProxy 是独立出口代理(http/https/socks5)，与记录流量的 MITM 代理无关。
|
||||
webSearchOn bool
|
||||
webSearchBackend string
|
||||
braveKey string
|
||||
tavilyKey string
|
||||
webSearchProxy string
|
||||
}
|
||||
|
||||
// Settings keys the UI toggles at runtime.
|
||||
const (
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
settingWorkers = "workers"
|
||||
// defaultWebSearchBackend is used when web search is on but no backend was picked.
|
||||
defaultWebSearchBackend = "ddgs"
|
||||
// defaultWorkers is the concurrent work-agent count when the setting is unset.
|
||||
defaultWorkers = 3
|
||||
)
|
||||
|
||||
// Workers returns the configured concurrent work-agent count (default 3). Read
|
||||
// per-task at engine.Run, so a change applies to tasks started afterwards.
|
||||
func (m *Manager) Workers() int {
|
||||
v, ok, err := m.pg.GetSetting(settingWorkers)
|
||||
if err != nil || !ok {
|
||||
return defaultWorkers
|
||||
}
|
||||
n, err := strconv.Atoi(strings.TrimSpace(v))
|
||||
if err != nil || n <= 0 {
|
||||
return defaultWorkers
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// SetWorkers persists the concurrent work-agent count. Values <=0 are rejected.
|
||||
func (m *Manager) SetWorkers(n int) error {
|
||||
if n <= 0 {
|
||||
return fmt.Errorf("workers 必须 >0")
|
||||
}
|
||||
return m.pg.SetSetting(settingWorkers, strconv.Itoa(n))
|
||||
}
|
||||
|
||||
// Enrich returns the asset auto-completion engine (may be nil if init failed).
|
||||
func (m *Manager) Enrich() *enrich.Engine { return m.enrich }
|
||||
|
||||
// NewManager connects to PostgreSQL and, if proxyAddr is non-empty, starts the
|
||||
// traffic-recording proxy. PostgreSQL is required (it is the single data source).
|
||||
func NewManager(dir, proxyAddr string) (*Manager, error) {
|
||||
// Resolve the data dir to an ABSOLUTE path up front. Every data path derives
|
||||
// from it — notably the MITM CA cert, whose path is injected into worker shells
|
||||
// (SSL_CERT_FILE/CURL_CA_BUNDLE) and read by WebFetch. A relative path (the
|
||||
// default is "./data" under `go run`) only resolves when the current working
|
||||
// directory happens to match, so curl/WebFetch in a different CWD fail to load
|
||||
// the CA → TLS to the proxy breaks (curl 000 / EOF). Absolute makes it CWD-proof.
|
||||
if abs, err := filepath.Abs(dir); err == nil {
|
||||
dir = abs
|
||||
}
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dsn, source, err := pgdb.DSN()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
log.Printf("[pg] 数据库配置来源: %s", source)
|
||||
pg, err := pgdb.Open(dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)}
|
||||
if proxyAddr != "" {
|
||||
tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr)
|
||||
if err != nil {
|
||||
log.Printf("[traffic] disabled: %v", err)
|
||||
} else {
|
||||
m.traffic = tr
|
||||
go func() {
|
||||
log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr())
|
||||
if err := tr.Start(); err != nil {
|
||||
log.Printf("[traffic] proxy stopped: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
// Asset auto-completion engine (§5): HTTP probes routed through the recording
|
||||
// proxy (via m.ProxyAddr, which honors the traffic-capture toggle).
|
||||
m.trafficOn = pg.GetBool(settingTrafficCapture, false)
|
||||
// Load persisted web-search config (default: off, ddgs).
|
||||
m.webSearchOn = pg.GetBool(settingWebSearchOn, false)
|
||||
if v, ok, _ := pg.GetSetting(settingWebSearchBackend); ok && v != "" {
|
||||
m.webSearchBackend = v
|
||||
} else {
|
||||
m.webSearchBackend = defaultWebSearchBackend
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingBraveKey); ok {
|
||||
m.braveKey = v
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingTavilyKey); ok {
|
||||
m.tavilyKey = v
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingWebSearchProxy); ok {
|
||||
m.webSearchProxy = v
|
||||
}
|
||||
m.enrich = enrich.New(m.assets, m.ProxyAddr, 4)
|
||||
// Reconcile the seeded browser MCP with the persisted capture state, so a
|
||||
// restart with capture already on keeps Playwright routed through the proxy.
|
||||
m.syncBrowserMCPProxy()
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// TrafficEnabled reports whether traffic capture is on (default off). When off,
|
||||
// no proxy/traffic tools/prompt are injected into agents (nothing is recorded).
|
||||
func (m *Manager) TrafficEnabled() bool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.trafficOn
|
||||
}
|
||||
|
||||
// SetTrafficEnabled persists and applies the traffic-capture toggle. Callers must
|
||||
// rebuild the agents (applyLLM) afterwards so the new proxy/tools/prompt take hold.
|
||||
func (m *Manager) SetTrafficEnabled(on bool) error {
|
||||
if err := m.pg.SetBool(settingTrafficCapture, on); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.trafficOn = on
|
||||
m.mu.Unlock()
|
||||
// Inject (on) or strip (off) the recording proxy + CA on the browser MCP so
|
||||
// Playwright routes through the MITM. Must run after the flag flip above, since
|
||||
// ProxyAddr/ProxyCACert honor it. putSettings rebuilds agents next (applyLLM),
|
||||
// which re-spawns the MCP with the new args/env.
|
||||
m.syncBrowserMCPProxy()
|
||||
return nil
|
||||
}
|
||||
|
||||
// WebSearch returns the current web-search config: whether it is enabled, the
|
||||
// backend ("ddgs" | "brave-free" | "tavily"), the Brave API key, the Tavily API
|
||||
// key (each empty unless set), and the dedicated egress proxy (empty = direct).
|
||||
func (m *Manager) WebSearch() (on bool, backend, braveKey, tavilyKey, proxy string) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
backend = m.webSearchBackend
|
||||
if backend == "" {
|
||||
backend = defaultWebSearchBackend
|
||||
}
|
||||
return m.webSearchOn, backend, m.braveKey, m.tavilyKey, m.webSearchProxy
|
||||
}
|
||||
|
||||
// WebSearchOpts returns the config as the agent-package struct the server pushes
|
||||
// into each agent. Disabled when off, or when a keyed backend is selected without
|
||||
// its key (so a half-configured backend never silently drops the tool at session build).
|
||||
func (m *Manager) WebSearchOpts() agent.WebSearchOpts {
|
||||
on, backend, braveKey, tavilyKey, proxy := m.WebSearch()
|
||||
if on && backend == "brave-free" && strings.TrimSpace(braveKey) == "" {
|
||||
on = false
|
||||
}
|
||||
if on && backend == "tavily" && strings.TrimSpace(tavilyKey) == "" {
|
||||
on = false
|
||||
}
|
||||
return agent.WebSearchOpts{Enabled: on, Backend: backend, BraveKey: braveKey, TavilyKey: tavilyKey, Proxy: proxy}
|
||||
}
|
||||
|
||||
// SetWebSearch persists and applies the web-search settings. braveKey, tavilyKey, and
|
||||
// proxy are each left untouched when nil (so toggling the switch doesn't wipe a saved
|
||||
// key/proxy; pass a pointer to "" to clear). Callers must rebuild agents (applyLLM)
|
||||
// afterwards so the settings take effect.
|
||||
func (m *Manager) SetWebSearch(on bool, backend string, braveKey, tavilyKey, proxy *string) error {
|
||||
backend = strings.TrimSpace(backend)
|
||||
if backend == "" {
|
||||
backend = defaultWebSearchBackend
|
||||
}
|
||||
if err := m.pg.SetBool(settingWebSearchOn, on); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.SetSetting(settingWebSearchBackend, backend); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.webSearchOn = on
|
||||
m.webSearchBackend = backend
|
||||
m.mu.Unlock()
|
||||
if braveKey != nil {
|
||||
if err := m.pg.SetSetting(settingBraveKey, *braveKey); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.braveKey = *braveKey
|
||||
m.mu.Unlock()
|
||||
}
|
||||
if tavilyKey != nil {
|
||||
if err := m.pg.SetSetting(settingTavilyKey, *tavilyKey); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.tavilyKey = *tavilyKey
|
||||
m.mu.Unlock()
|
||||
}
|
||||
if proxy != nil {
|
||||
p := strings.TrimSpace(*proxy)
|
||||
if err := m.pg.SetSetting(settingWebSearchProxy, p); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.webSearchProxy = p
|
||||
m.mu.Unlock()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// browserMCPName is the seeded Playwright MCP whose proxy args + CA env are kept
|
||||
// in sync with the traffic-capture toggle.
|
||||
const browserMCPName = "browser"
|
||||
|
||||
// syncBrowserMCPProxy reconciles the seeded browser MCP's proxy args + CA env with
|
||||
// the current traffic-capture state: capture on → route Playwright through the
|
||||
// recording proxy (--proxy-server) and trust its MITM CA (NODE_EXTRA_CA_CERTS);
|
||||
// capture off → strip both. Idempotent, and a no-op if the user deleted/renamed the
|
||||
// MCP. Must be called WITHOUT m.mu held (ProxyAddr/ProxyCACert take the lock).
|
||||
func (m *Manager) syncBrowserMCPProxy() {
|
||||
servers, err := m.pg.ListMCP()
|
||||
if err != nil {
|
||||
log.Printf("[mcp] browser 代理同步: 读取 MCP 列表失败: %v", err)
|
||||
return
|
||||
}
|
||||
var srv *pgdb.MCPServer
|
||||
for _, s := range servers {
|
||||
if s.Name == browserMCPName {
|
||||
srv = s
|
||||
break
|
||||
}
|
||||
}
|
||||
if srv == nil {
|
||||
return // user removed/renamed it — leave it alone
|
||||
}
|
||||
|
||||
proxy := m.ProxyAddr() // "" when capture off
|
||||
cert := m.ProxyCACert() // "" when capture off
|
||||
|
||||
args := stripProxyArgs(decodeStrSlice(srv.Args))
|
||||
env := decodeStrMap(srv.Env)
|
||||
delete(env, "NODE_EXTRA_CA_CERTS")
|
||||
if proxy != "" {
|
||||
args = append(args, "--proxy-server", proxy)
|
||||
if cert != "" {
|
||||
env["NODE_EXTRA_CA_CERTS"] = cert
|
||||
}
|
||||
}
|
||||
srv.Args = encodeJSON(args)
|
||||
srv.Env = encodeJSON(env)
|
||||
if _, err := m.pg.SaveMCP(srv); err != nil {
|
||||
log.Printf("[mcp] browser 代理同步失败: %v", err)
|
||||
return
|
||||
}
|
||||
if proxy != "" {
|
||||
log.Printf("[mcp] browser MCP 已挂捕获代理 %s (CA %s)", proxy, cert)
|
||||
} else {
|
||||
log.Printf("[mcp] browser MCP 已移除捕获代理配置")
|
||||
}
|
||||
}
|
||||
|
||||
// stripProxyArgs removes any --proxy-server/--proxy-bypass flags (both "--flag val"
|
||||
// and "--flag=val" forms) so they can be re-added cleanly from current state,
|
||||
// without mutating the input slice.
|
||||
func stripProxyArgs(args []string) []string {
|
||||
out := make([]string, 0, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
a := args[i]
|
||||
if a == "--proxy-server" || a == "--proxy-bypass" {
|
||||
i++ // skip the following value too
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(a, "--proxy-server=") || strings.HasPrefix(a, "--proxy-bypass=") {
|
||||
continue
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeStrSlice(raw json.RawMessage) []string {
|
||||
var out []string
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeStrMap(raw json.RawMessage) map[string]string {
|
||||
out := map[string]string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeJSON(v any) json.RawMessage {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// HostTools are runtime host-provided tools added to EVERY agent's base list (via
|
||||
// ToolAugment); the tools table then filters them per-agent binding. Currently the
|
||||
// traffic tools, gated by the global capture switch: empty when capture is off, so
|
||||
// no agent gets traffic_search/traffic_get regardless of binding.
|
||||
func (m *Manager) HostTools() []actool.CoreTool {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return nil
|
||||
}
|
||||
return m.traffic.Tools()
|
||||
}
|
||||
|
||||
func (m *Manager) Assets() *pgdb.AssetStore { return m.assets }
|
||||
func (m *Manager) PG() *pgdb.DB { return m.pg }
|
||||
func (m *Manager) Traffic() *traffic.Traffic { return m.traffic }
|
||||
|
||||
// ProxyAddr returns the recording proxy address agents route through — empty when
|
||||
// traffic capture is off, so no proxy is injected (agent runs direct, no recording).
|
||||
func (m *Manager) ProxyAddr() string {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return ""
|
||||
}
|
||||
return m.traffic.ProxyAddr()
|
||||
}
|
||||
|
||||
// ProxyCACert returns the recording proxy's CA cert path (empty when no proxy or
|
||||
// traffic capture is off), which WebFetch trusts to verify HTTPS through the MITM.
|
||||
func (m *Manager) ProxyCACert() string {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return ""
|
||||
}
|
||||
return m.traffic.CACertPath()
|
||||
}
|
||||
|
||||
func (m *Manager) Close() error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.traffic != nil {
|
||||
m.traffic.Close()
|
||||
}
|
||||
return m.pg.Close()
|
||||
}
|
||||
|
||||
// isTerminalStatus reports whether a task status is terminal (done/failed/timeout).
|
||||
// Package-local shim over db.IsTerminal so all server files share one definition.
|
||||
func isTerminalStatus(status string) bool { return pgdb.IsTerminal(status) }
|
||||
|
||||
// unixOrZero returns t's unix seconds, or 0 when the time is nil.
|
||||
func unixOrZero(t *time.Time) int64 {
|
||||
if t == nil {
|
||||
return 0
|
||||
}
|
||||
return t.Unix()
|
||||
}
|
||||
|
||||
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,
|
||||
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
|
||||
LLMProfileID: pt.LLMProfileID,
|
||||
TimeoutSeconds: pt.TimeoutSeconds, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
|
||||
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
|
||||
}
|
||||
}
|
||||
|
||||
// CreateTask creates a task + its exploration and makes it active.
|
||||
// timeoutSeconds is the task-level wall-clock budget (0 = 不限时).
|
||||
func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds int) (*Task, error) {
|
||||
pt, err := m.pg.CreateTask(description, goal, llmProfileID, timeoutSeconds)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
t := taskFromPG(pt, m.pg.Exploration(pt.ExplorationID), m.interceptor)
|
||||
m.mu.Lock()
|
||||
m.tasks[t.ID] = t
|
||||
m.active = t.ID
|
||||
m.mu.Unlock()
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// LoadExisting rebuilds in-memory task handles from the PG task registry.
|
||||
func (m *Manager) LoadExisting() []*Task {
|
||||
pts, err := m.pg.ListTasks()
|
||||
if err != nil {
|
||||
log.Printf("[manager] reload: %v", err)
|
||||
return nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
var loaded []*Task
|
||||
for _, pt := range pts {
|
||||
id := strconv.FormatInt(pt.ID, 10)
|
||||
if _, ok := m.tasks[id]; ok {
|
||||
continue
|
||||
}
|
||||
t := taskFromPG(pt, m.pg.Exploration(pt.ExplorationID), m.interceptor)
|
||||
m.tasks[id] = t
|
||||
loaded = append(loaded, t)
|
||||
}
|
||||
if m.active == "" {
|
||||
var newest *Task
|
||||
for _, t := range m.tasks {
|
||||
if newest == nil || t.CreatedAt > newest.CreatedAt {
|
||||
newest = t
|
||||
}
|
||||
}
|
||||
if newest != nil {
|
||||
m.active = newest.ID
|
||||
}
|
||||
}
|
||||
if len(loaded) > 0 {
|
||||
log.Printf("[manager] reloaded %d task(s) from PG", len(loaded))
|
||||
}
|
||||
return loaded
|
||||
}
|
||||
|
||||
// SetTaskPaused persists a task's paused state.
|
||||
func (m *Manager) SetTaskPaused(id string, paused bool) error {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return m.pg.SetPaused(n, paused)
|
||||
}
|
||||
|
||||
// TaskStatus returns a task's current in-memory status (empty if unknown).
|
||||
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 ""
|
||||
}
|
||||
|
||||
// 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) {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
m.mu.RLock()
|
||||
timeout := 0
|
||||
if t := m.tasks[id]; t != nil {
|
||||
timeout = t.TimeoutSeconds
|
||||
}
|
||||
m.mu.RUnlock()
|
||||
dl, err := m.pg.StampFirstRun(n, timeout)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var dlUnix int64
|
||||
if dl != nil {
|
||||
dlUnix = dl.Unix()
|
||||
}
|
||||
m.mu.Lock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
if t.FirstRunAt == 0 {
|
||||
t.FirstRunAt = time.Now().Unix()
|
||||
}
|
||||
t.DeadlineAt = dlUnix
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return dlUnix, nil
|
||||
}
|
||||
|
||||
// SetTaskStatusGuarded sets a TERMINAL status only if the task isn't already terminal
|
||||
// (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) {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
won, err = m.pg.SetTerminalStatusGuarded(n, status)
|
||||
if err != nil || !won {
|
||||
return won, err
|
||||
}
|
||||
m.mu.Lock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
t.Status = status
|
||||
if t.CompletedAt == 0 {
|
||||
t.CompletedAt = time.Now().Unix()
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// 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 {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.SetStatus(n, status); err != nil {
|
||||
return err
|
||||
}
|
||||
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()
|
||||
}
|
||||
} else {
|
||||
t.CompletedAt = 0
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteTask removes a task (and its exploration subgraph) from PG + memory.
|
||||
func (m *Manager) DeleteTask(id string) error {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.DeleteTask(n); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
delete(m.tasks, id)
|
||||
if m.active == id {
|
||||
m.active = ""
|
||||
for _, t := range m.tasks {
|
||||
m.active = t.ID
|
||||
break
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) Task(id string) (*Task, bool) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
t, ok := m.tasks[id]
|
||||
return t, ok
|
||||
}
|
||||
|
||||
// ActiveTask returns the currently active task (or nil).
|
||||
func (m *Manager) ActiveTask() *Task {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
if m.active == "" {
|
||||
return nil
|
||||
}
|
||||
return m.tasks[m.active]
|
||||
}
|
||||
|
||||
// SetActive switches the active task. Returns false if the id is unknown.
|
||||
func (m *Manager) SetActive(id string) bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if _, ok := m.tasks[id]; !ok {
|
||||
return false
|
||||
}
|
||||
m.active = id
|
||||
return true
|
||||
}
|
||||
|
||||
func (m *Manager) ResolveTask(id string) *Task {
|
||||
if id == "" || id == "active" {
|
||||
return m.ActiveTask()
|
||||
}
|
||||
t, _ := m.Task(id)
|
||||
return t
|
||||
}
|
||||
|
||||
func (m *Manager) List() []*Task {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
out := make([]*Task, 0, len(m.tasks))
|
||||
for _, t := range m.tasks {
|
||||
out = append(out, t)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
|
||||
return out
|
||||
}
|
||||
|
||||
// Notify signals that the asset/exploration graph changed (debounced consumer
|
||||
// wakes the planner). Non-blocking.
|
||||
func (t *Task) Notify() {
|
||||
select {
|
||||
case t.notify <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// NotifyDone is Notify plus a hint: intentID just finished and is what triggered
|
||||
// this wake-up. The planner reads the accumulated ids next round so it knows
|
||||
// which intents' fresh yields to focus on. Ids pile up (debounce) until the
|
||||
// round drains them via drainDone.
|
||||
func (t *Task) NotifyDone(intentID int64) {
|
||||
if intentID > 0 {
|
||||
t.doneMu.Lock()
|
||||
t.doneIntents = append(t.doneIntents, intentID)
|
||||
t.doneMu.Unlock()
|
||||
}
|
||||
t.Notify()
|
||||
}
|
||||
|
||||
// drainDone returns and clears the intent ids completed since the last round.
|
||||
func (t *Task) drainDone() []int64 {
|
||||
t.doneMu.Lock()
|
||||
defer t.doneMu.Unlock()
|
||||
ids := t.doneIntents
|
||||
t.doneIntents = nil
|
||||
return ids
|
||||
}
|
||||
Reference in new issue
Block a user