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,362 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/norma/skill"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// wireAgentAugment connects the PG agent_visibility table into the agent runtime:
|
||||
// an agent's visible skills are loaded from the filesystem and packed into one
|
||||
// Skill meta-tool; its visible stdio MCP servers are spawned and expanded to
|
||||
// mcp__server__tool. skillDir is the root directory of all skill subdirectories.
|
||||
// hostTools, if set, returns runtime host tools (currently the traffic tools when
|
||||
// capture is on) to add to EVERY agent's base list — the DB tools table then
|
||||
// filters them per-agent binding. Empty/nil → no host tools this run (capture off).
|
||||
func wireAgentAugment(pg *db.DB, skillDir string, hostTools func() ([]actool.CoreTool, map[string][]string)) {
|
||||
agent.ToolAugment = func(ctx context.Context, agentKey string) ([]actool.CoreTool, agent.DeferredInfo, func()) {
|
||||
a, err := pg.GetAgentByKey(agentKey)
|
||||
if err != nil || a == nil {
|
||||
return nil, agent.DeferredInfo{}, nil
|
||||
}
|
||||
var extra []actool.CoreTool
|
||||
|
||||
// --- skills: load visible skills into reg (used for the Skill meta-tool
|
||||
// and to know which MCP servers are skill-gated). ---
|
||||
var reg *skill.Registry
|
||||
if names, _ := pg.AgentSkillNames(a.ID); len(names) > 0 {
|
||||
nameSet := make(map[string]bool, len(names))
|
||||
for _, n := range names {
|
||||
nameSet[n] = true
|
||||
}
|
||||
if allReg, err := skill.LoadDir(skillDir); err == nil && allReg != nil {
|
||||
reg = skill.NewRegistry()
|
||||
for _, s := range allReg.List() {
|
||||
// match by directory name (Base of Dir), not by skill display Name
|
||||
if s.Dir != "" && nameSet[filepath.Base(s.Dir)] {
|
||||
reg.Add(s)
|
||||
}
|
||||
}
|
||||
if len(reg.List()) == 0 {
|
||||
reg = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
// A server named by any visible skill's `mcps:` is skill-gated: its tools
|
||||
// are deferred + locked (not in the global block) until that skill loads.
|
||||
gated := map[string]bool{}
|
||||
if reg != nil {
|
||||
for _, s := range reg.List() {
|
||||
for _, srv := range s.MCPs {
|
||||
gated[srv] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- mcp: connect enabled servers that are directly visible to this agent
|
||||
// OR skill-gated (named in a visible skill's mcps field). Directly-visible
|
||||
// and NOT gated → global (unlocked from session start). Skill-gated →
|
||||
// deferred until that skill is invoked (regardless of direct visibility).
|
||||
var closers []io.Closer
|
||||
serverTools := map[string][]string{} // server name → its tool names
|
||||
var allNames, globalNames []string
|
||||
globalSet := map[string]bool{}
|
||||
{
|
||||
mcpIDs, _ := pg.AgentVisible(a.ID, "mcp")
|
||||
want := idSet(mcpIDs)
|
||||
all, _ := pg.ListMCP()
|
||||
for _, m := range all {
|
||||
if !m.Enabled {
|
||||
continue
|
||||
}
|
||||
directVisible := want[m.ID]
|
||||
skillGated := gated[m.Name]
|
||||
if !directVisible && !skillGated {
|
||||
continue // neither directly visible nor referenced by a visible skill
|
||||
}
|
||||
cl, err := connectMCP(ctx, m)
|
||||
if err != nil {
|
||||
log.Printf("[mcp] %s 连接失败: %v", m.Name, err)
|
||||
continue
|
||||
}
|
||||
closers = append(closers, cl)
|
||||
ts, err := cl.Tools(ctx)
|
||||
if err != nil {
|
||||
log.Printf("[mcp] %s tools/list 失败: %v", m.Name, err)
|
||||
continue
|
||||
}
|
||||
for _, t := range ts {
|
||||
extra = append(extra, t)
|
||||
allNames = append(allNames, t.Name())
|
||||
serverTools[m.Name] = append(serverTools[m.Name], t.Name())
|
||||
if directVisible && !skillGated {
|
||||
// directly visible and not gated → available from session start
|
||||
globalNames = append(globalNames, t.Name())
|
||||
globalSet[t.Name()] = true
|
||||
}
|
||||
// skill-gated tools stay out of globalNames; unlocked via unlockSkill()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Shared call-gate: global MCP tools are unlocked from the start; skill-gated
|
||||
// ones are unlocked when their skill loads (or replayed from history — C2).
|
||||
unlock := actool.NewUnlockSet(globalNames...)
|
||||
unlockSkill := func(skillName string) {
|
||||
if reg == nil {
|
||||
return
|
||||
}
|
||||
if s, ok := reg.Get(skillName); ok {
|
||||
for _, srv := range s.MCPs {
|
||||
unlock.Add(serverTools[srv]...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Skill meta-tool: on load, unlock the skill's MCPs and reveal their names
|
||||
// (the ones not already in the global block).
|
||||
if reg != nil {
|
||||
reg.OnInvoke = func(s skill.Skill) string {
|
||||
unlockSkill(s.Name)
|
||||
var reveal []string
|
||||
for _, srv := range s.MCPs {
|
||||
for _, n := range serverTools[srv] {
|
||||
if !globalSet[n] {
|
||||
reveal = append(reveal, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
return actool.RenderDeferredToolsBlock(reveal)
|
||||
}
|
||||
extra = append(extra, reg.Tool())
|
||||
}
|
||||
|
||||
// host tools (traffic / orchestration / custom) — added to every agent's base;
|
||||
// ToolResolve then keeps them only for agents the tool is bound to. Custom
|
||||
// tools flagged deferred contribute their names to the deferred wiring so
|
||||
// their schema is withheld (SearchExtraTools/ExecuteExtraTool), same as MCP.
|
||||
if hostTools != nil {
|
||||
ht, deferredBinds := hostTools()
|
||||
extra = append(extra, ht...)
|
||||
for name, boundAgents := range deferredBinds {
|
||||
if !contains(boundAgents, a.Key) {
|
||||
continue // only defer names this agent is actually bound to
|
||||
}
|
||||
allNames = append(allNames, name) // schema withheld from the prompt
|
||||
globalNames = append(globalNames, name) // advertised in the deferred block
|
||||
globalSet[name] = true
|
||||
if unlock != nil {
|
||||
unlock.Add(name) // global deferred → callable from the start
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
for _, c := range closers {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
def := agent.DeferredInfo{
|
||||
Deferred: allNames,
|
||||
GlobalNames: globalNames,
|
||||
Unlock: unlock,
|
||||
UnlockSkill: unlockSkill,
|
||||
}
|
||||
return extra, def, cleanup
|
||||
}
|
||||
}
|
||||
|
||||
func idSet(ids []int64) map[int64]bool {
|
||||
m := make(map[int64]bool, len(ids))
|
||||
for _, id := range ids {
|
||||
m[id] = true
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// seedPrompts writes each built-in agent's code-default prompt body into
|
||||
// agent_prompts on startup — first-insert only (SeedPromptIfEmpty is a no-op once
|
||||
// any version exists), so the DB becomes the authoritative editable source while
|
||||
// user edits survive restarts. Runs after seedBuiltins has created the agent rows.
|
||||
func seedPrompts(pg *db.DB) {
|
||||
for key, tmpl := range agent.BuiltinPromptSeeds() {
|
||||
a, err := pg.GetAgentByKey(key)
|
||||
if err != nil || a == nil {
|
||||
log.Printf("[prompts] seed %s 跳过: agent 不存在 (%v)", key, err)
|
||||
continue
|
||||
}
|
||||
if err := pg.SeedPromptIfEmpty(a.ID, tmpl); err != nil {
|
||||
log.Printf("[prompts] seed %s 失败: %v", key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wireTools seeds the built-in tool catalog (idempotent, first-insert only so page
|
||||
// edits survive restart) and wires the DB tools table into the agent runtime: at
|
||||
// tool-assembly time each built-in tool is filtered by its agent binding / enabled
|
||||
// flag and, if kept, wrapped so the model sees the DB-overridden description/schema
|
||||
// and缺省入参 get injected. MCP/skill/host tools have no row and pass through.
|
||||
func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) {
|
||||
// Seed the built-in domain tools (first-insert only; DO NOTHING preserves edits).
|
||||
// No startup prune: rows we didn't seed are left alone so future user-defined
|
||||
// custom tools (system=false, added via the UI) survive restarts.
|
||||
for _, s := range agent.BuiltinToolSeeds() {
|
||||
schema, _ := json.Marshal(s.Schema)
|
||||
agents, _ := json.Marshal(s.Agents)
|
||||
if err := pg.SeedTool(s.Key, s.Desc, schema, agents); err != nil {
|
||||
log.Printf("[tools] seed %s 失败: %v", s.Key, err)
|
||||
}
|
||||
}
|
||||
// Seed the traffic host tools so they're bindable per-agent like built-ins.
|
||||
// Default binding = worker (preserves prior behavior). Their runtime availability
|
||||
// is still gated by the global capture switch (hostTools() returns them only when
|
||||
// capture is on), so an off-capture binding simply never surfaces the tool.
|
||||
trafficAgents, _ := json.Marshal([]string{"worker"})
|
||||
for _, t := range traffic.SeedToolMetas() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
if err := pg.SeedTool(t.Name(), t.Description(), schema, trafficAgents); err != nil {
|
||||
log.Printf("[tools] seed %s 失败: %v", t.Name(), err)
|
||||
}
|
||||
}
|
||||
// bashInteractiveShellNote is appended to Bash's description ONLY for agents whose
|
||||
// interactive_shell is on, so Bash points at shell_open for interactive programs
|
||||
// without ever referencing a tool that isn't injected (§14.1/§14.2).
|
||||
const bashInteractiveShellNote = "\n\n需要【交互输入】的程序（msfconsole / ssh 交互登录 / mysql、psql、python 等 REPL / 密码或 yes/no 提示 / nc 反弹 shell）不要用 Bash（它没有 stdin、会卡住），改用 shell_open 开交互会话（用完 shell_close）。一次性、非交互命令仍用 Bash。"
|
||||
agent.ToolResolve = func(_ context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool {
|
||||
rows, err := pg.ListTools()
|
||||
if err != nil {
|
||||
log.Printf("[tools] 读取工具表失败，按代码默认放行: %v", err)
|
||||
return tools
|
||||
}
|
||||
byKey := make(map[string]*db.Tool, len(rows))
|
||||
for _, t := range rows {
|
||||
byKey[t.Key] = t
|
||||
}
|
||||
out := tools[:0:0]
|
||||
for _, t := range tools {
|
||||
row, known := byKey[t.Name()]
|
||||
if !known { // MCP/skill/host tool: no row → untouched
|
||||
out = append(out, t)
|
||||
continue
|
||||
}
|
||||
if !row.Enabled || !contains(row.Agents, agentKey) {
|
||||
continue // disabled globally or not bound to this agent → drop
|
||||
}
|
||||
var schema map[string]any
|
||||
if len(row.Schema) > 0 {
|
||||
_ = json.Unmarshal(row.Schema, &schema)
|
||||
}
|
||||
out = append(out, agent.DecorateTool(t, row.Description, schema))
|
||||
}
|
||||
// inject: domain tools bound to this agent in the DB but absent from the
|
||||
// incoming list. Covers agents (Auto, custom) whose base only has DefaultTools
|
||||
// and therefore never includes ToolSet-backed domain tools. Per-task instances
|
||||
// in the base always win: inList is built from the original incoming list so a
|
||||
// worker's own upsert_asset is never shadowed by the server-level registry copy.
|
||||
if len(domainReg) > 0 {
|
||||
inList := make(map[string]bool, len(tools))
|
||||
for _, t := range tools {
|
||||
inList[t.Name()] = true
|
||||
}
|
||||
for _, row := range rows {
|
||||
if row.Kind == "shell" || !row.Enabled || !contains(row.Agents, agentKey) || inList[row.Key] {
|
||||
continue
|
||||
}
|
||||
inst, ok := domainReg[row.Key]
|
||||
if !ok {
|
||||
continue // not a domain tool; custom/host tools are injected via hostTools()
|
||||
}
|
||||
var schema map[string]any
|
||||
if len(row.Schema) > 0 {
|
||||
_ = json.Unmarshal(row.Schema, &schema)
|
||||
}
|
||||
out = append(out, agent.DecorateTool(inst, row.Description, schema))
|
||||
}
|
||||
}
|
||||
// shell hints: user-defined kind="shell" tools are not callable — they are
|
||||
// environment declarations that tell the model which command-line tools are
|
||||
// installed. Collect the ones bound to this agent and append to Bash's description.
|
||||
var shellHints []string
|
||||
for _, row := range rows {
|
||||
if row.Kind == "shell" && row.Enabled && contains(row.Agents, agentKey) {
|
||||
shellHints = append(shellHints, "- "+row.Key+": "+row.Description)
|
||||
}
|
||||
}
|
||||
if len(shellHints) > 0 {
|
||||
note := "\n\n以下工具已安装在此 bash 环境中，可直接通过 Bash 调用：\n" + strings.Join(shellHints, "\n")
|
||||
for i, t := range out {
|
||||
if t.Name() == "Bash" {
|
||||
out[i] = agent.DecorateTool(t, t.Description()+note, t.InputSchema())
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
// interactive shell: gated purely by the agent's interactive_shell flag (like
|
||||
// web_search), NOT by tools-table binding. When on, inject the 5 shell_* tools
|
||||
// and COUPLE the Bash description addendum so it points at shell_open — and never
|
||||
// dangles when off. See docs/交互式shell设计.md §14.2.
|
||||
if !actool.InteractiveShellDisabled() {
|
||||
if a, err := pg.GetAgentByKey(agentKey); err == nil && a != nil && a.InteractiveShell {
|
||||
out = append(out, actool.ShellSessionTools()...)
|
||||
for i, t := range out {
|
||||
if t.Name() == "Bash" {
|
||||
out[i] = agent.DecorateTool(t, t.Description()+bashInteractiveShellNote, t.InputSchema())
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
|
||||
func contains(ss []string, v string) bool {
|
||||
for _, s := range ss {
|
||||
if s == v {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// buildDomainReg builds a name→CoreTool registry from a server-level ToolSet
|
||||
// (real AssetStore, nil ExplorationStore, taskID=0). Used by ToolResolve to inject
|
||||
// domain tools into agents (Auto, custom) that don't own a per-task ToolSet.
|
||||
// nil as → returns nil (no injection, graceful degradation).
|
||||
func buildDomainReg(as *db.AssetStore) map[string]actool.CoreTool {
|
||||
if as == nil {
|
||||
return nil
|
||||
}
|
||||
serverTS := agent.NewToolSet(nil, "")
|
||||
serverTS.SetAssetStore(as, as.Companies())
|
||||
reg := make(map[string]actool.CoreTool)
|
||||
for _, t := range serverTS.AllDomainTools() {
|
||||
reg[t.Name()] = t
|
||||
}
|
||||
return reg
|
||||
}
|
||||
|
||||
func jsonStrSlice(raw json.RawMessage) []string {
|
||||
var out []string
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func jsonStrMap(raw json.RawMessage) map[string]string {
|
||||
out := map[string]string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// hasSkillTool reports whether the packed tool set contains the Skill meta-tool.
|
||||
func hasSkillTool(tools []actool.CoreTool) bool {
|
||||
for _, t := range tools {
|
||||
if t.Name() == "Skill" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestAssembleVisibleSkill verifies a filesystem skill made visible to an agent is
|
||||
// packed into that agent's tool set (as the single Skill meta-tool) via the wired
|
||||
// ToolAugment hook. Skills live on disk under skillDir; visibility is a per-agent
|
||||
// (agent × skill_name) row keyed by the skill's directory name.
|
||||
func TestAssembleVisibleSkill(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v) — skipping", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
|
||||
// Filesystem skill fixture: <skillDir>/t-assemble/SKILL.md. The directory name
|
||||
// (t-assemble) is the visibility key matched against AgentSkillNames.
|
||||
skillDir := t.TempDir()
|
||||
const skillName = "t-assemble"
|
||||
if err := os.MkdirAll(filepath.Join(skillDir, skillName), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
md := "---\nname: t-assemble\ndescription: do the thing\n---\ndo the thing"
|
||||
if err := os.WriteFile(filepath.Join(skillDir, skillName, "SKILL.md"), []byte(md), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
wireAgentAugment(pg, skillDir, nil)
|
||||
t.Cleanup(func() { agent.ToolAugment = nil })
|
||||
|
||||
ag, _ := pg.GetAgentByKey("planner")
|
||||
if ag == nil {
|
||||
t.Fatal("planner agent missing")
|
||||
}
|
||||
|
||||
// Isolate this test from whatever skill visibility the planner already has in the
|
||||
// shared DB: clear it now, restore it on cleanup.
|
||||
orig, _ := pg.AgentSkillNames(ag.ID)
|
||||
if err := pg.SetAgentSkillVisibility(ag.ID, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = pg.SetAgentSkillVisibility(ag.ID, orig) })
|
||||
|
||||
// before: skill not visible → no Skill meta-tool
|
||||
extra, _, cleanup := agent.ToolAugment(context.Background(), "planner")
|
||||
cleanup()
|
||||
if hasSkillTool(extra) {
|
||||
t.Fatalf("expected no Skill meta-tool before the skill is made visible")
|
||||
}
|
||||
|
||||
// make the skill visible to planner
|
||||
if err := pg.ToggleSkillVisibility(ag.ID, skillName, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = pg.ToggleSkillVisibility(ag.ID, skillName, false) })
|
||||
|
||||
// after: the Skill meta-tool is packed in
|
||||
extra, _, cleanup = agent.ToolAugment(context.Background(), "planner")
|
||||
defer cleanup()
|
||||
if !hasSkillTool(extra) {
|
||||
t.Fatalf("expected the Skill meta-tool after making the skill visible")
|
||||
}
|
||||
|
||||
// AugmentTools appends extra to base without filtering base
|
||||
combined, _, c2 := agent.AugmentTools(context.Background(), "planner", nil)
|
||||
defer c2()
|
||||
if !hasSkillTool(combined) {
|
||||
t.Fatalf("AugmentTools should surface the Skill meta-tool")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,426 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// assetStore returns the asset store.
|
||||
func (s *Server) assetStore() *db.AssetStore {
|
||||
if s.m.pg == nil {
|
||||
return nil
|
||||
}
|
||||
return s.m.pg.Assets()
|
||||
}
|
||||
|
||||
// companyStore returns the company store.
|
||||
func (s *Server) companyStore() *db.CompanyStore {
|
||||
if s.m.pg == nil {
|
||||
return nil
|
||||
}
|
||||
return s.m.pg.Companies()
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/companies
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) listCompanies(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
companies, err := cs.ListCompanies()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, companies)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) createCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
Logo string `json:"logo"`
|
||||
Scope []string `json:"scope"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
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())
|
||||
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
|
||||
}
|
||||
}
|
||||
writeJSON(w, 201, out)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/companies/{id}
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) getCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid id")
|
||||
return
|
||||
}
|
||||
c, err := cs.GetCompany(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if c == nil {
|
||||
writeErr(w, 404, "company not found")
|
||||
return
|
||||
}
|
||||
scope, err := cs.GetScope(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"company": c, "scope": scope})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies/{id}/scope
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) addCompanyScope(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid company id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Scope []string `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())
|
||||
return
|
||||
}
|
||||
var added, skipped, invalid int
|
||||
var errs []string
|
||||
if req.Reset {
|
||||
added, invalid, errs = cs.UpdateScope(id, req.Scope, req.Reason)
|
||||
} else {
|
||||
added, skipped, invalid, errs = cs.AddScope(id, req.Scope, req.Reason)
|
||||
}
|
||||
out := map[string]any{
|
||||
"added": added,
|
||||
"skipped": skipped,
|
||||
"invalid": invalid,
|
||||
}
|
||||
if len(errs) > 0 {
|
||||
out["errors"] = errs
|
||||
}
|
||||
writeJSON(w, 200, out)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// DELETE /api/companies/{id}
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) deleteCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid id")
|
||||
return
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := cs.DeleteCompany(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": 1, "assets_deleted": assetsDeleted})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies/reattribute
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
if err := cs.RecomputeAttribution(); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
q := r.URL.Query()
|
||||
typ := q.Get("type")
|
||||
limit, _ := strconv.Atoi(q.Get("limit"))
|
||||
offset, _ := strconv.Atoi(q.Get("offset"))
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
var 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 := 0
|
||||
|
||||
if dsl := q.Get("dsl"); dsl != "" {
|
||||
assets, err = as.QueryDSL(dsl, typ, limit, offset)
|
||||
if err == nil {
|
||||
total, _ = as.CountDSL(dsl, typ)
|
||||
}
|
||||
} 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)
|
||||
total = len(assets)
|
||||
case taskID > 0:
|
||||
assets, err = as.QueryByTask(taskID, typ, limit)
|
||||
total = len(assets)
|
||||
default:
|
||||
if typ == "" {
|
||||
typ = "root_domain"
|
||||
}
|
||||
assets, err = as.QueryByType(typ, limit, offset)
|
||||
if err == nil {
|
||||
total, _ = as.CountByType(typ)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"count": len(assets),
|
||||
"total": total,
|
||||
"assets": assets,
|
||||
})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/assets/counts
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) assetCounts(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
counts, err := as.CountsByType()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, counts)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// DELETE /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) deleteAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
IDs []int64 `json:"ids"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if len(req.IDs) == 0 {
|
||||
writeErr(w, 400, "ids required")
|
||||
return
|
||||
}
|
||||
deleted, err := as.DeleteByIDs(req.IDs)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": deleted})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
TaskID int64 `json:"task_id"`
|
||||
Assets []struct {
|
||||
Type string `json:"type"`
|
||||
// root_domain / subdomain
|
||||
Domain string `json:"domain"`
|
||||
ICP string `json:"icp"`
|
||||
RecordType string `json:"record_type"`
|
||||
RecordValue []string `json:"record_value"`
|
||||
// ip
|
||||
IP string `json:"ip"`
|
||||
BoundDomains []string `json:"bound_domains"`
|
||||
OpenPorts []db.PortService `json:"open_ports"`
|
||||
// app
|
||||
AppName string `json:"app_name"`
|
||||
BundleID string `json:"bundle_id"`
|
||||
Category string `json:"category"`
|
||||
Description string `json:"description"`
|
||||
AppICP string `json:"app_icp"`
|
||||
// service http
|
||||
URL string `json:"url"`
|
||||
Technologies []string `json:"technologies"`
|
||||
StatusCode *int `json:"status_code"`
|
||||
ContentLength *int64 `json:"content_length"`
|
||||
PageTitle string `json:"page_title"`
|
||||
FaviconMMH3 string `json:"favicon_mmh3"`
|
||||
Auth []map[string]any `json:"auth"`
|
||||
ServiceIP string `json:"service_ip"`
|
||||
// service other
|
||||
Port int `json:"port"`
|
||||
ServiceName string `json:"service_name"`
|
||||
// endpoint
|
||||
Method string `json:"method"`
|
||||
Params []map[string]any `json:"params"`
|
||||
} `json:"assets"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
type result struct {
|
||||
Index int `json:"index"`
|
||||
ID int64 `json:"id"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
type errEntry struct {
|
||||
Index int `json:"index"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
var results []result
|
||||
var errs []errEntry
|
||||
|
||||
for i, a := range req.Assets {
|
||||
var id int64
|
||||
var err error
|
||||
switch a.Type {
|
||||
case "root_domain":
|
||||
id, err = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: a.Domain, ICP: a.ICP, TaskID: req.TaskID})
|
||||
case "ip":
|
||||
id, err = as.UpsertIP(db.UpsertIPReq{IP: a.IP, BoundDomains: a.BoundDomains, OpenPorts: a.OpenPorts, TaskID: req.TaskID})
|
||||
case "subdomain":
|
||||
id, err = as.UpsertSubdomain(db.UpsertSubdomainReq{Domain: a.Domain, RecordType: a.RecordType, RecordValue: a.RecordValue, ICP: a.ICP, TaskID: req.TaskID})
|
||||
case "app":
|
||||
id, err = as.UpsertApp(db.UpsertAppReq{Name: a.AppName, BundleID: a.BundleID, Category: a.Category, Description: a.Description, ICP: a.AppICP, TaskID: req.TaskID})
|
||||
case "service":
|
||||
if a.URL != "" {
|
||||
svcIP := a.ServiceIP
|
||||
if svcIP == "" {
|
||||
svcIP = a.IP
|
||||
}
|
||||
id, err = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: a.URL, Technologies: a.Technologies, StatusCode: a.StatusCode, ContentLength: a.ContentLength, PageTitle: a.PageTitle, FaviconMMH3: a.FaviconMMH3, Auth: a.Auth, IP: svcIP, TaskID: req.TaskID})
|
||||
} else {
|
||||
id, err = as.UpsertOtherService(db.UpsertOtherServiceReq{Domain: a.Domain, IP: a.IP, Port: a.Port, ServiceName: a.ServiceName, Auth: a.Auth, TaskID: req.TaskID})
|
||||
}
|
||||
case "endpoint":
|
||||
svcIP := a.ServiceIP
|
||||
if svcIP == "" {
|
||||
svcIP = a.IP
|
||||
}
|
||||
id, err = as.UpsertEndpoint(db.UpsertEndpointReq{URL: a.URL, Method: a.Method, Params: a.Params, IP: svcIP, TaskID: req.TaskID})
|
||||
default:
|
||||
errs = append(errs, errEntry{Index: i, Error: "unknown type: " + a.Type})
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
errs = append(errs, errEntry{Index: i, Error: err.Error()})
|
||||
continue
|
||||
}
|
||||
results = append(results, result{Index: i, ID: id, Type: a.Type})
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"results": results, "errors": errs})
|
||||
}
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"log"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
const (
|
||||
jwtKeyFilename = "jwt.key"
|
||||
authPassKey = "auth.password_hash"
|
||||
jwtTTL = 7 * 24 * time.Hour
|
||||
keyChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
)
|
||||
|
||||
// loadOrCreateJWTKey reads the 32-byte signing key from dataDir/jwt.key.
|
||||
// On first run it generates a cryptographically random key and persists it.
|
||||
func loadOrCreateJWTKey(dataDir string) ([]byte, error) {
|
||||
path := filepath.Join(dataDir, jwtKeyFilename)
|
||||
if data, err := os.ReadFile(path); err == nil && len(strings.TrimSpace(string(data))) >= 32 {
|
||||
return []byte(strings.TrimSpace(string(data))), nil
|
||||
}
|
||||
buf := make([]byte, 32)
|
||||
for i := range buf {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(int64(len(keyChars))))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate jwt key: %w", err)
|
||||
}
|
||||
buf[i] = keyChars[n.Int64()]
|
||||
}
|
||||
if err := os.WriteFile(path, buf, 0600); err != nil {
|
||||
return nil, fmt.Errorf("write jwt key: %w", err)
|
||||
}
|
||||
log.Printf("[auth] 新 JWT key 已写入 %s", path)
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// signJWT issues a 7-day HS256 token for user ARTEX.
|
||||
func signJWT(key []byte) (string, error) {
|
||||
return jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.RegisteredClaims{
|
||||
Subject: "ARTEX",
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(jwtTTL)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
}).SignedString(key)
|
||||
}
|
||||
|
||||
// verifyJWT returns true when tokenStr is a valid, non-expired HS256 token.
|
||||
func verifyJWT(tokenStr string, key []byte) bool {
|
||||
t, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) {
|
||||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method")
|
||||
}
|
||||
return key, nil
|
||||
})
|
||||
return err == nil && t.Valid
|
||||
}
|
||||
|
||||
// extractToken reads the JWT from Authorization: Bearer header,
|
||||
// artex_token cookie, or ?token= query param (for SSE connections).
|
||||
func extractToken(r *http.Request) string {
|
||||
if h := r.Header.Get("Authorization"); strings.HasPrefix(h, "Bearer ") {
|
||||
return strings.TrimPrefix(h, "Bearer ")
|
||||
}
|
||||
if c, err := r.Cookie("artex_token"); err == nil && c.Value != "" {
|
||||
return c.Value
|
||||
}
|
||||
return r.URL.Query().Get("token")
|
||||
}
|
||||
|
||||
// requireAuth wraps h with JWT validation.
|
||||
// /api/auth/* and /api/health are exempt.
|
||||
func (s *Server) requireAuth(h http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
p := r.URL.Path
|
||||
if strings.HasPrefix(p, "/api/auth/") || p == "/api/health" {
|
||||
h.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
tok := extractToken(r)
|
||||
if tok == "" {
|
||||
writeErr(w, 401, "未授权")
|
||||
return
|
||||
}
|
||||
if !verifyJWT(tok, s.jwtKey) {
|
||||
writeErr(w, 401, "token 无效或已过期")
|
||||
return
|
||||
}
|
||||
h.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// GET /api/auth/status — reports whether the admin password has been initialised.
|
||||
func (s *Server) authStatus(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
hash, _, _ := pg.GetSetting(authPassKey)
|
||||
writeJSON(w, 200, map[string]any{"initialized": hash != ""})
|
||||
}
|
||||
|
||||
// POST /api/auth/init — sets the password for the first time; rejected if already set.
|
||||
func (s *Server) authInit(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
existing, _, _ := pg.GetSetting(authPassKey)
|
||||
if existing != "" {
|
||||
writeErr(w, 403, "密码已设置")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil || req.Password == "" {
|
||||
writeErr(w, 400, "密码不能为空")
|
||||
return
|
||||
}
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcrypt.DefaultCost)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "密码加密失败")
|
||||
return
|
||||
}
|
||||
if err := pg.SetSetting(authPassKey, string(hash)); err != nil {
|
||||
writeErr(w, 500, "保存失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "token 生成失败")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"token": tok})
|
||||
}
|
||||
|
||||
// POST /api/auth/change-password — changes the admin password. Requires a valid
|
||||
// token (this route is under /api/auth/* which requireAuth exempts, so the token
|
||||
// is validated here) AND the current password.
|
||||
func (s *Server) authChangePassword(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
if !verifyJWT(extractToken(r), s.jwtKey) {
|
||||
writeErr(w, 401, "未授权")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
OldPassword string `json:"old_password"`
|
||||
NewPassword string `json:"new_password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "请求格式错误")
|
||||
return
|
||||
}
|
||||
if req.NewPassword == "" {
|
||||
writeErr(w, 400, "新密码不能为空")
|
||||
return
|
||||
}
|
||||
hash, ok, _ := pg.GetSetting(authPassKey)
|
||||
if !ok || hash == "" {
|
||||
writeErr(w, 403, "密码未初始化，请先设置密码")
|
||||
return
|
||||
}
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.OldPassword)); err != nil {
|
||||
writeErr(w, 401, "当前密码错误")
|
||||
return
|
||||
}
|
||||
newHash, err := bcrypt.GenerateFromPassword([]byte(req.NewPassword), bcrypt.DefaultCost)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "密码加密失败")
|
||||
return
|
||||
}
|
||||
if err := pg.SetSetting(authPassKey, string(newHash)); err != nil {
|
||||
writeErr(w, 500, "保存失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// POST /api/auth/login — validates username/password and returns a JWT.
|
||||
func (s *Server) authLogin(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "请求格式错误")
|
||||
return
|
||||
}
|
||||
if req.Username != "ARTEX" {
|
||||
writeErr(w, 401, "用户名或密码错误")
|
||||
return
|
||||
}
|
||||
hash, ok, _ := pg.GetSetting(authPassKey)
|
||||
if !ok || hash == "" {
|
||||
writeErr(w, 403, "密码未初始化，请先设置密码")
|
||||
return
|
||||
}
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.Password)); err != nil {
|
||||
writeErr(w, 401, "用户名或密码错误")
|
||||
return
|
||||
}
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "token 生成失败")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"token": tok})
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Broadcaster is a per-task in-process pub/sub for live activity events. The
|
||||
// engine publishes each appended activity at its single emit point; SSE handlers
|
||||
// subscribe per task. Storage (activity table) and the live stream come from the
|
||||
// same Publish call, so they never diverge.
|
||||
type Broadcaster struct {
|
||||
mu sync.Mutex
|
||||
subs map[string]map[chan db.Activity]struct{} // task id -> set of subscriber channels
|
||||
}
|
||||
|
||||
func NewBroadcaster() *Broadcaster {
|
||||
return &Broadcaster{subs: map[string]map[chan db.Activity]struct{}{}}
|
||||
}
|
||||
|
||||
// Subscribe returns a buffered channel of activities for a task plus an
|
||||
// unsubscribe func the caller must invoke (defer) to release it.
|
||||
func (b *Broadcaster) Subscribe(task string) (<-chan db.Activity, func()) {
|
||||
ch := make(chan db.Activity, 256)
|
||||
b.mu.Lock()
|
||||
if b.subs[task] == nil {
|
||||
b.subs[task] = map[chan db.Activity]struct{}{}
|
||||
}
|
||||
b.subs[task][ch] = struct{}{}
|
||||
b.mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return ch, func() {
|
||||
once.Do(func() {
|
||||
b.mu.Lock()
|
||||
if m := b.subs[task]; m != nil {
|
||||
delete(m, ch)
|
||||
if len(m) == 0 {
|
||||
delete(b.subs, task)
|
||||
}
|
||||
}
|
||||
b.mu.Unlock()
|
||||
close(ch)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Publish fans an activity out to all subscribers of a task. Non-blocking: if a
|
||||
// subscriber's buffer is full the event is dropped — the client reconnects with
|
||||
// its last seq cursor and catches up the gap from the DB, so liveness never
|
||||
// stalls the engine.
|
||||
func (b *Broadcaster) Publish(task string, a db.Activity) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
for ch := range b.subs[task] {
|
||||
select {
|
||||
case ch <- a:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
// ---------- conversations (chat page) ----------
|
||||
//
|
||||
// A conversation is a ChatGPT-style thread bound to an agent key, independent of
|
||||
// the pentest task graph. Turns run on ChatAgent; steps persist to
|
||||
// conversation_activities and the browser POLLS ?since=cursor for live updates
|
||||
// (no per-conversation SSE broadcaster needed).
|
||||
|
||||
func (s *Server) pgListConversations(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
cs, err := pg.ListConversations()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"conversations": cs})
|
||||
}
|
||||
|
||||
func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
AgentKey string `json:"agent_key"`
|
||||
Title string `json:"title"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
req.AgentKey = strings.TrimSpace(req.AgentKey)
|
||||
if req.AgentKey == "" {
|
||||
writeErr(w, 400, "agent_key 不能为空")
|
||||
return
|
||||
}
|
||||
a, err := pg.GetAgentByKey(req.AgentKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if a == nil {
|
||||
writeErr(w, 404, "agent 不存在")
|
||||
return
|
||||
}
|
||||
if req.LLMProfileID != nil {
|
||||
if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok {
|
||||
writeErr(w, 400, "指定的 LLM 配置不存在或未设置 API Key")
|
||||
return
|
||||
}
|
||||
}
|
||||
title := strings.TrimSpace(req.Title)
|
||||
if title == "" {
|
||||
title = "新对话"
|
||||
}
|
||||
c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, c)
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
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())
|
||||
return
|
||||
}
|
||||
if req.LLMProfileID != nil {
|
||||
if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok {
|
||||
writeErr(w, 400, "指定的 LLM 配置不存在或未设置 API Key")
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := pg.UpdateConversationProfile(c.ID, req.LLMProfileID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// convByID resolves the {id} path value to a conversation (404 if absent).
|
||||
func (s *Server) convByID(w http.ResponseWriter, r *http.Request) (*db.DB, *db.Conversation, bool) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return nil, nil, false
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad conversation id")
|
||||
return nil, nil, false
|
||||
}
|
||||
c, err := pg.GetConversation(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return nil, nil, false
|
||||
}
|
||||
if c == nil {
|
||||
writeErr(w, 404, "conversation not found")
|
||||
return nil, nil, false
|
||||
}
|
||||
return pg, c, true
|
||||
}
|
||||
|
||||
func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req struct{ Title string }
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
title := strings.TrimSpace(req.Title)
|
||||
if title == "" {
|
||||
writeErr(w, 400, "标题不能为空")
|
||||
return
|
||||
}
|
||||
if err := pg.RenameConversation(c.ID, title); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteConversation(c.ID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": c.ID})
|
||||
}
|
||||
|
||||
func (s *Server) convBusyKey(id int64) string { return "conv-" + strconv.FormatInt(id, 10) }
|
||||
|
||||
// pgConversationMessages returns steps after ?since=cursor plus whether a turn is
|
||||
// still running (so the client knows to keep polling). Same item shape as the task
|
||||
// activity stream, so the frontend transcript renderer is reused verbatim.
|
||||
func (s *Server) pgConversationMessages(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
since, _ := strconv.ParseInt(r.URL.Query().Get("since"), 10, 64)
|
||||
items, cursor, err := pg.ConvActivityList(c.ID, since, 0)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.chatMu.Lock()
|
||||
running := s.chatBusy[s.convBusyKey(c.ID)]
|
||||
s.chatMu.Unlock()
|
||||
writeJSON(w, 200, map[string]any{"items": activityDTOs(items), "cursor": cursor, "running": running})
|
||||
}
|
||||
|
||||
// pgConversationMsgDetail lazily returns one step's full detail blob.
|
||||
func (s *Server) pgConversationMsgDetail(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
seq, ok := pathInt(r, "seq")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad seq")
|
||||
return
|
||||
}
|
||||
detail, err := pg.ConvActivityDetail(c.ID, seq)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"detail": detail})
|
||||
}
|
||||
|
||||
// pgStopConversation aborts the in-flight run for one conversation (manual stop —
|
||||
// the run/stop button in the chat UI). It cancels only THIS session's agent run;
|
||||
// the P3 trigger queue is untouched, so the agent's next queued fire still starts.
|
||||
func (s *Server) pgStopConversation(w http.ResponseWriter, r *http.Request) {
|
||||
_, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
cancel := s.chatCancel[busyKey]
|
||||
s.chatMu.Unlock()
|
||||
if cancel == nil {
|
||||
writeJSON(w, 200, map[string]any{"status": "idle"}) // nothing running
|
||||
return
|
||||
}
|
||||
cancel()
|
||||
writeJSON(w, 200, map[string]any{"status": "stopping"})
|
||||
}
|
||||
|
||||
// pgSendConversationMessage persists the human turn, then runs the agent in the
|
||||
// background (steps stream to conversation_activities, polled by the client). One
|
||||
// in-flight turn per conversation.
|
||||
func (s *Server) pgSendConversationMessage(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req struct{ Message string }
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
msg := strings.TrimSpace(req.Message)
|
||||
if msg == "" {
|
||||
writeErr(w, 400, "消息不能为空")
|
||||
return
|
||||
}
|
||||
ca := s.chatAgentRef()
|
||||
if ca == nil {
|
||||
writeErr(w, 503, "LLM 未配置，无法对话")
|
||||
return
|
||||
}
|
||||
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
if s.chatBusy[busyKey] {
|
||||
s.chatMu.Unlock()
|
||||
writeErr(w, 409, "该会话正在处理上一条消息，请稍候")
|
||||
return
|
||||
}
|
||||
s.chatBusy[busyKey] = true
|
||||
s.chatMu.Unlock()
|
||||
|
||||
// persist + float the human turn; auto-title the thread from the first message.
|
||||
// Worker is the agent key (not "user") so the transcript stays a single lane
|
||||
// (no worker chips) — the kind='user' already right-aligns it as a human bubble.
|
||||
if _, err := pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "user", Summary: firstLine(msg, 200), Detail: msg}); err != nil {
|
||||
log.Printf("[conv %d] append user msg failed: %v", c.ID, err)
|
||||
}
|
||||
if c.Title == "" || c.Title == "新对话" {
|
||||
_ = pg.RenameConversation(c.ID, firstLine(msg, 40))
|
||||
}
|
||||
_ = pg.TouchConversation(c.ID)
|
||||
|
||||
s.runConversation(c, msg, busyKey)
|
||||
writeJSON(w, 202, map[string]any{"status": "accepted"})
|
||||
}
|
||||
|
||||
// runConversation runs ONE agent turn on a conversation in a background goroutine
|
||||
// (each conversation is independent → parallel). Steps stream to
|
||||
// conversation_activities. Shared by the chat HTTP handler and the P3 scheduler.
|
||||
// busyKey clears when the run ends (best-effort in-flight marker).
|
||||
func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) {
|
||||
go s.runConversationSync(c, msg, busyKey)
|
||||
}
|
||||
|
||||
// runConversationSync runs ONE agent turn on a conversation and BLOCKS until the
|
||||
// run ends (clearing busyKey). runConversation wraps it in a goroutine for the
|
||||
// fire-and-forget chat path; the P3 trigger queue calls it directly so it can wait
|
||||
// for completion before starting the next queued fire for the same agent.
|
||||
func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) {
|
||||
// Per-run cancellable context so a manual stop (pgStopConversation) can abort
|
||||
// just this session. Registered under chatMu so the stop handler can find it.
|
||||
ctx, cancel := context.WithCancel(intercept.WithConvID(s.ctx, c.ID))
|
||||
s.chatMu.Lock()
|
||||
s.chatCancel[busyKey] = cancel
|
||||
s.chatMu.Unlock()
|
||||
defer func() {
|
||||
cancel()
|
||||
s.chatMu.Lock()
|
||||
delete(s.chatBusy, busyKey)
|
||||
delete(s.chatCancel, busyKey)
|
||||
s.chatMu.Unlock()
|
||||
}()
|
||||
ca := s.chatAgentRef()
|
||||
if c.LLMProfileID != nil {
|
||||
if pa := s.chatAgentForProfile(*c.LLMProfileID); pa != nil {
|
||||
ca = pa
|
||||
}
|
||||
}
|
||||
if ca == nil {
|
||||
return
|
||||
}
|
||||
pg := s.m.pg
|
||||
maxTurns := s.agentMaxTurns(c.AgentKey)
|
||||
maxDuration := time.Duration(s.agentRunSeconds(c.AgentKey)) * time.Second
|
||||
webSearch := false
|
||||
if a, err := pg.GetAgentByKey(c.AgentKey); err == nil && a != nil {
|
||||
webSearch = a.WebSearch
|
||||
}
|
||||
sessionID := s.convBusyKey(c.ID) // "conv-<id>" transcript session
|
||||
emit := func(rec db.Activity) {
|
||||
if _, err := pg.AppendConvActivity(c.ID, rec); err != nil {
|
||||
log.Printf("[conv %d] append activity failed: %v", c.ID, err)
|
||||
}
|
||||
}
|
||||
// On a manual stop ctx is cancelled; Chat already emits a clean "已手动停止"
|
||||
// step, so skip the raw-error entry — only surface genuine failures.
|
||||
if _, err := ca.Chat(ctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil && ctx.Err() == nil {
|
||||
_, _ = pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "text", IsError: true,
|
||||
Summary: "（出错：" + err.Error() + "）", Detail: err.Error()})
|
||||
}
|
||||
_ = pg.TouchConversation(c.ID)
|
||||
}
|
||||
|
||||
// StartTriggeredRun enqueues a P3 trigger fire for agentKey. Fires for the SAME
|
||||
// agent queue up and run one at a time (FIFO) — the next only starts after the
|
||||
// previous conversation finishes. Distinct agents still run concurrently. taskID +
|
||||
// mergeable let the drainer coalesce same-task event triggers before a run starts.
|
||||
func (s *Server) StartTriggeredRun(agentKey, title, message string, taskID int64, mergeable bool) {
|
||||
if s.m.pg == nil || s.chatAgentRef() == nil {
|
||||
return
|
||||
}
|
||||
s.queueMu.Lock()
|
||||
s.triggerQ[agentKey] = append(s.triggerQ[agentKey], triggeredRun{agentKey: agentKey, title: title, message: message, taskID: taskID, mergeable: mergeable})
|
||||
depth := len(s.triggerQ[agentKey])
|
||||
if !s.triggerRun[agentKey] {
|
||||
s.triggerRun[agentKey] = true
|
||||
go s.drainTriggerQueue(agentKey)
|
||||
} else {
|
||||
log.Printf("[trigger] %s busy → queued (depth %d)", agentKey, depth)
|
||||
}
|
||||
s.queueMu.Unlock()
|
||||
}
|
||||
|
||||
// drainTriggerQueue runs one agent's queued trigger fires serially until the queue
|
||||
// empties, then clears the active flag. Before each run it coalesces all queued
|
||||
// finding/goal fires from the SAME task into one conversation (nextTriggerRun), so a
|
||||
// burst of a task's events becomes a single session. A panic is contained.
|
||||
func (s *Server) drainTriggerQueue(agentKey string) {
|
||||
for {
|
||||
s.queueMu.Lock()
|
||||
if len(s.triggerQ[agentKey]) == 0 || s.ctx.Err() != nil {
|
||||
delete(s.triggerQ, agentKey)
|
||||
delete(s.triggerRun, agentKey)
|
||||
s.queueMu.Unlock()
|
||||
return
|
||||
}
|
||||
item := s.nextTriggerRun(agentKey)
|
||||
s.queueMu.Unlock()
|
||||
|
||||
s.runTriggeredRun(item)
|
||||
}
|
||||
}
|
||||
|
||||
// nextTriggerRun pops the head of the agent's queue; if the head is a mergeable
|
||||
// (finding/goal) fire, it also pulls every other queued mergeable fire from the SAME
|
||||
// task and merges them into one run. Caller holds queueMu.
|
||||
func (s *Server) nextTriggerRun(agentKey string) triggeredRun {
|
||||
q := s.triggerQ[agentKey]
|
||||
head := q[0]
|
||||
if !head.mergeable || head.taskID == 0 {
|
||||
s.triggerQ[agentKey] = q[1:]
|
||||
return head
|
||||
}
|
||||
group := []triggeredRun{head}
|
||||
rest := q[:0:0] // keep non-matching items in order
|
||||
for _, it := range q[1:] {
|
||||
if it.mergeable && it.taskID == head.taskID {
|
||||
group = append(group, it)
|
||||
} else {
|
||||
rest = append(rest, it)
|
||||
}
|
||||
}
|
||||
s.triggerQ[agentKey] = rest
|
||||
return mergeTriggeredRuns(group)
|
||||
}
|
||||
|
||||
// mergeTriggeredRuns folds several same-task event fires into one run: a header plus
|
||||
// each fire's message, so the agent handles the task's burst in a single conversation.
|
||||
func mergeTriggeredRuns(items []triggeredRun) triggeredRun {
|
||||
if len(items) == 1 {
|
||||
return items[0]
|
||||
}
|
||||
first := items[0]
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "【本会话合并了任务 #%d 的 %d 条触发事件，请一并处理】\n", first.taskID, len(items))
|
||||
for i, it := range items {
|
||||
fmt.Fprintf(&b, "\n── 触发 %d ──\n%s\n", i+1, it.message)
|
||||
}
|
||||
return triggeredRun{
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("合并触发 · task#%d · %d 条", first.taskID, len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
}
|
||||
}
|
||||
|
||||
// runTriggeredRun creates the conversation for one queued fire, records the human
|
||||
// turn, and runs it synchronously (blocks until the run ends). recover keeps a
|
||||
// panic from killing the drain loop and wedging the agent's queue.
|
||||
func (s *Server) runTriggeredRun(item triggeredRun) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Printf("[trigger] run for %s panicked: %v", item.agentKey, r)
|
||||
}
|
||||
}()
|
||||
pg := s.m.pg
|
||||
c, err := pg.CreateConversation(item.agentKey, firstLine(item.title, 60), nil)
|
||||
if err != nil {
|
||||
log.Printf("[trigger] create conversation for %s failed: %v", item.agentKey, err)
|
||||
return
|
||||
}
|
||||
if _, err := pg.AppendConvActivity(c.ID, db.Activity{Worker: item.agentKey, Kind: "user", Summary: firstLine(item.message, 200), Detail: item.message}); err != nil {
|
||||
log.Printf("[trigger] append msg failed: %v", err)
|
||||
}
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
s.chatBusy[busyKey] = true
|
||||
s.chatMu.Unlock()
|
||||
s.runConversationSync(c, item.message, busyKey)
|
||||
}
|
||||
|
||||
// firstLine returns a single-line, length-capped preview (shared with summaries).
|
||||
func firstLine(s string, max int) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if i := strings.IndexAny(s, "\r\n"); i >= 0 {
|
||||
s = s[:i]
|
||||
}
|
||||
if len([]rune(s)) > max {
|
||||
s = string([]rune(s)[:max]) + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// TestCoreTaskLifecyclePG exercises the migrated core (tasks/exploration on PG)
|
||||
// through the real HTTP mux: create → goal nodes seeded → list → delete cascade.
|
||||
func TestCoreTaskLifecyclePG(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)
|
||||
h := s.Handler()
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatalf("sign jwt: %v", err)
|
||||
}
|
||||
|
||||
do := func(method, path string, body any) (int, map[string]any) {
|
||||
var r *http.Request
|
||||
if body != nil {
|
||||
b, _ := json.Marshal(body)
|
||||
r = httptest.NewRequest(method, path, bytes.NewReader(b))
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
r.Header.Set("Authorization", "Bearer "+tok)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, r)
|
||||
var out map[string]any
|
||||
json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
// doRetry retries on 5xx (transient DB conflicts from parallel test packages).
|
||||
doRetry := func(method, path string, body any) (int, map[string]any) {
|
||||
var code int
|
||||
var out map[string]any
|
||||
for i := 0; i < 5; i++ {
|
||||
code, out = do(method, path, body)
|
||||
if code < 500 {
|
||||
break
|
||||
}
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
}
|
||||
return code, out
|
||||
}
|
||||
|
||||
// create a task → 201, returns PG task (string id + exploration_id)
|
||||
code, out := doRetry("POST", "/api/tasks", map[string]string{"description": "smoke", "goal": "测试 SQLi/XSS"})
|
||||
if code != 201 {
|
||||
t.Fatalf("create task: %d (%v)", code, out)
|
||||
}
|
||||
id, _ := out["id"].(string)
|
||||
expID := int64(out["exploration_id"].(float64))
|
||||
if id == "" || expID == 0 {
|
||||
t.Fatalf("bad task payload: %v", out)
|
||||
}
|
||||
|
||||
// the exploration owns goal node(s) — goal seeding is async, poll briefly
|
||||
var goals []*db.Node
|
||||
for i := 0; i < 30; i++ {
|
||||
goals, err = m.pg.Exploration(expID).ListByKind("goal", 10)
|
||||
if err != nil || len(goals) > 0 {
|
||||
break
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(goals) < 1 {
|
||||
t.Fatalf("expected goal nodes seeded, got %d", len(goals))
|
||||
}
|
||||
|
||||
// it shows in the task list
|
||||
code, out = doRetry("GET", "/api/tasks", nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("list tasks: %d", code)
|
||||
}
|
||||
|
||||
// delete → cascade removes the exploration subgraph
|
||||
code, _ = doRetry("DELETE", "/api/tasks/"+id, nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("delete task: %d", code)
|
||||
}
|
||||
nid, _ := strconv.ParseInt(id, 10, 64)
|
||||
if got, _ := m.pg.GetTask(nid); got != nil {
|
||||
t.Fatalf("task should be deleted from PG")
|
||||
}
|
||||
var n int
|
||||
m.pg.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, expID).Scan(&n)
|
||||
if n != 0 {
|
||||
t.Fatalf("exploration nodes should be cascade-deleted, got %d", n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,543 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 本文件实现自定义工具执行器(docs/自定义工具设计.md)。system=false 的 tools 行按
|
||||
// kind 分派:command(渲染命令→复用 Bash 底层 run)、script(仅 Python;写临时文件、
|
||||
// stdin=参数 JSON + env TOOL_*、用配置的解释器)、http(原生请求+可设代理)。这些工具
|
||||
// 像流量/编排工具一样 seed 不需要(它们本就在 tools 表),经 hostTools 注入、按绑定过滤。
|
||||
|
||||
// ---------- 自定义工具 CRUD ----------
|
||||
|
||||
type customToolReq struct {
|
||||
Key string `json:"key"`
|
||||
Description string `json:"description"`
|
||||
Schema json.RawMessage `json:"schema"`
|
||||
Agents []string `json:"agents"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Kind string `json:"kind"` // command | script | http
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Deferred bool `json:"deferred"`
|
||||
}
|
||||
|
||||
var reToolKey = reAgentKey // 同 agent key 规则:小写字母开头 + 小写字母/数字/下划线
|
||||
|
||||
func (s *Server) pgCreateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req customToolReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
req.Key = strings.TrimSpace(req.Key)
|
||||
if !reToolKey.MatchString(req.Key) {
|
||||
writeErr(w, 400, "key 需小写字母开头，仅含小写字母/数字/下划线")
|
||||
return
|
||||
}
|
||||
if req.Kind != "command" && req.Kind != "script" && req.Kind != "http" && req.Kind != "shell" {
|
||||
writeErr(w, 400, "kind 需为 command / script / http / shell")
|
||||
return
|
||||
}
|
||||
if req.Kind == "http" && !hasSchemaProps(req.Schema) {
|
||||
writeErr(w, 400, "http 工具必须提供参数 JSON Schema(不能留空)")
|
||||
return
|
||||
}
|
||||
if exist, _ := pg.GetTool(req.Key); exist != nil {
|
||||
writeErr(w, 409, "该 key 已存在(内置或自定义工具)")
|
||||
return
|
||||
}
|
||||
if err := pg.CreateCustomTool(&db.Tool{
|
||||
Key: req.Key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"key": req.Key})
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
key := r.PathValue("key")
|
||||
existing, err := pg.GetTool(key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if existing == nil || existing.System {
|
||||
writeErr(w, 400, "只能编辑自定义工具")
|
||||
return
|
||||
}
|
||||
var req customToolReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.Kind != "command" && req.Kind != "script" && req.Kind != "http" && req.Kind != "shell" {
|
||||
writeErr(w, 400, "kind 需为 command / script / http / shell")
|
||||
return
|
||||
}
|
||||
if req.Kind == "http" && !hasSchemaProps(req.Schema) {
|
||||
writeErr(w, 400, "http 工具必须提供参数 JSON Schema(不能留空)")
|
||||
return
|
||||
}
|
||||
if err := pg.UpdateCustomTool(&db.Tool{
|
||||
Key: key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
key := r.PathValue("key")
|
||||
if err := pg.DeleteCustomTool(key); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": key})
|
||||
}
|
||||
|
||||
// testToolReq is a dry-run request from the editor: run the given (possibly unsaved)
|
||||
// exec spec with sample params, without persisting the tool. Same executor path as a
|
||||
// real tool call — it runs arbitrary command/script/http on the server, which the
|
||||
// custom-tool feature already allows, so no new capability is granted.
|
||||
type testToolReq struct {
|
||||
Kind string `json:"kind"` // command | script | http
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Params map[string]any `json:"params"`
|
||||
}
|
||||
|
||||
// pgTestCustomTool executes an exec spec once and returns its raw output + error
|
||||
// flag, so the editor can debug a tool before saving it. Per-kind timeouts still
|
||||
// apply from the exec spec (with defaults); the outer ceiling is a hard backstop.
|
||||
func (s *Server) pgTestCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req testToolReq
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "无效的请求体")
|
||||
return
|
||||
}
|
||||
params := req.Params
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
tc := &actool.ToolContext{WorkingDir: s.m.dir} // run in the project dir, like a real call
|
||||
var res actool.Result
|
||||
switch req.Kind {
|
||||
case "command":
|
||||
res, _ = s.runCommandTool(ctx, req.Exec, params, tc)
|
||||
case "script":
|
||||
res, _ = s.runScriptTool(ctx, "test", req.Exec, params, tc)
|
||||
case "http":
|
||||
res, _ = s.runHTTPTool(ctx, req.Exec, params)
|
||||
case "shell":
|
||||
writeErr(w, 400, "shell 类型工具是 bash 环境声明，无可执行内容")
|
||||
return
|
||||
default:
|
||||
writeErr(w, 400, "未知工具类型: "+req.Kind)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"output": res.Flatten(), "is_error": res.IsError})
|
||||
}
|
||||
|
||||
// ---------- Python 解释器(检测 + 入库 + 覆盖) ----------
|
||||
|
||||
const settingPythonInterp = "python_interpreter"
|
||||
|
||||
// detectPython finds a python interpreter absolute path (python3 preferred).
|
||||
func detectPython() string {
|
||||
for _, c := range []string{"python3", "python"} {
|
||||
if p, err := exec.LookPath(c); err == nil {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// pythonInterpreter resolves the interpreter: user-set > stored auto-detect > live
|
||||
// detect. "" only when truly none found.
|
||||
func (s *Server) pythonInterpreter() string {
|
||||
if v, ok, _ := s.m.pg.GetSetting(settingPythonInterp); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return detectPython()
|
||||
}
|
||||
|
||||
// seedPythonInterpreter stores the auto-detected interpreter on startup if unset
|
||||
// (never clobbers a user-set value).
|
||||
func (s *Server) seedPythonInterpreter() {
|
||||
if v, ok, _ := s.m.pg.GetSetting(settingPythonInterp); ok && strings.TrimSpace(v) != "" {
|
||||
return
|
||||
}
|
||||
if p := detectPython(); p != "" {
|
||||
_ = s.m.pg.SetSetting(settingPythonInterp, p)
|
||||
log.Printf("[custom-tool] 自动检测到 python 解释器: %s", p)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- exec 规格 ----------
|
||||
|
||||
type commandExec struct {
|
||||
Command string `json:"command"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
}
|
||||
type scriptExec struct {
|
||||
Code string `json:"code"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
}
|
||||
type httpExec struct {
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
Body string `json:"body"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
Proxy string `json:"proxy"`
|
||||
UseRecordingProxy bool `json:"use_recording_proxy"`
|
||||
}
|
||||
|
||||
func timeoutOr(ms, def int) time.Duration {
|
||||
if ms <= 0 {
|
||||
return time.Duration(def) * time.Millisecond
|
||||
}
|
||||
return time.Duration(ms) * time.Millisecond
|
||||
}
|
||||
|
||||
// ---------- 通用工具构造 ----------
|
||||
|
||||
// customTools builds CoreTools for every user-defined (system=false) tool row.
|
||||
// shell-kind tools are environment hints only — they surface in the Bash tool
|
||||
// description via ToolResolve and do NOT create callable tool entries here.
|
||||
func (s *Server) customTools() ([]actool.CoreTool, error) {
|
||||
rows, err := s.m.pg.ListCustomTools()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]actool.CoreTool, 0, len(rows))
|
||||
for _, t := range rows {
|
||||
if t.Kind == "shell" {
|
||||
continue // shell hints are handled by ToolResolve → Bash description
|
||||
}
|
||||
out = append(out, s.buildCustomTool(t))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildCustomTool turns one custom-tool row into a CoreTool. Empty schema → a thin
|
||||
// {args:string} (薄壳工具), so command/http templates can use {args}.
|
||||
func (s *Server) buildCustomTool(t *db.Tool) actool.CoreTool {
|
||||
schema := ensureSchema(t.Schema)
|
||||
key, kind, execRaw := t.Key, t.Kind, t.Exec
|
||||
run := func(ctx context.Context, in json.RawMessage, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var params map[string]any
|
||||
if len(in) > 0 {
|
||||
_ = json.Unmarshal(in, ¶ms)
|
||||
}
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
switch kind {
|
||||
case "command":
|
||||
return s.runCommandTool(ctx, execRaw, params, tc)
|
||||
case "script":
|
||||
return s.runScriptTool(ctx, key, execRaw, params, tc)
|
||||
case "http":
|
||||
return s.runHTTPTool(ctx, execRaw, params)
|
||||
default:
|
||||
return actool.Errorf("未知自定义工具类型: " + kind), nil
|
||||
}
|
||||
}
|
||||
return actool.Build(actool.Spec{
|
||||
Name: key, Description: t.Description, Schema: schema,
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: run,
|
||||
})
|
||||
}
|
||||
|
||||
// hasSchemaProps reports whether raw is a JSON-Schema object with ≥1 property.
|
||||
// http tools require an explicit schema (the auto {args} shell can't name the
|
||||
// {param} placeholders in URL/headers/body), so an empty schema is rejected.
|
||||
func hasSchemaProps(raw json.RawMessage) bool {
|
||||
if len(raw) == 0 {
|
||||
return false
|
||||
}
|
||||
var m map[string]any
|
||||
if json.Unmarshal(raw, &m) != nil {
|
||||
return false
|
||||
}
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
return len(props) > 0
|
||||
}
|
||||
|
||||
// ensureSchema returns the tool's schema, or a thin {args:string} when none given.
|
||||
func ensureSchema(raw json.RawMessage) map[string]any {
|
||||
var m map[string]any
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &m)
|
||||
}
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
if len(props) > 0 {
|
||||
return m
|
||||
}
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"args": map[string]any{"type": "string", "description": "命令/参数(自由文本)"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- command:渲染命令 → 复用 Bash 底层 run ----------
|
||||
|
||||
func (s *Server) runCommandTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec commandExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
if strings.TrimSpace(spec.Command) == "" {
|
||||
return actool.Errorf("command 为空"), nil
|
||||
}
|
||||
cmd := renderTemplate(spec.Command, params, shellQuote)
|
||||
// 复用 Bash 也在用的底层 run(经 Bash CoreTool.Call):自动继承安全 floor/超时/
|
||||
// 代理 env/输出溢出。工具与 Bash 平级、共用底层,不经过 Bash 这个工具让模型调。
|
||||
bashIn, _ := json.Marshal(map[string]any{"command": cmd})
|
||||
if spec.TimeoutMs > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, timeoutOr(spec.TimeoutMs, 120000))
|
||||
defer cancel()
|
||||
}
|
||||
return actool.NewBash().Call(ctx, bashIn, tc)
|
||||
}
|
||||
|
||||
// ---------- script(仅 Python):临时文件 + stdin JSON + env ----------
|
||||
|
||||
func (s *Server) runScriptTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec scriptExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
if strings.TrimSpace(spec.Code) == "" {
|
||||
return actool.Errorf("script code 为空"), nil
|
||||
}
|
||||
interp := s.pythonInterpreter()
|
||||
if interp == "" {
|
||||
return actool.Errorf("未配置且未检测到 python 解释器(在系统配置里设置)"), nil
|
||||
}
|
||||
workDir := s.m.dir
|
||||
var sessionEnv []string
|
||||
if tc != nil {
|
||||
if tc.WorkingDir != "" {
|
||||
workDir = tc.WorkingDir
|
||||
}
|
||||
sessionEnv = tc.Env
|
||||
}
|
||||
body, err := execPython(ctx, interp, key, spec.Code, params, workDir, sessionEnv, timeoutOr(spec.TimeoutMs, 120000))
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(clipOutput(body, tc)), nil
|
||||
}
|
||||
|
||||
// execPython writes the code to a temp .py under workDir/.tools, runs it via interp
|
||||
// with the params JSON on stdin + scalar params mirrored to TOOL_<NAME> env, and
|
||||
// returns combined stdout+stderr (with a timeout/exit note). Standalone + testable.
|
||||
func execPython(ctx context.Context, interp, key, code string, params map[string]any, workDir string, sessionEnv []string, timeout time.Duration) (string, error) {
|
||||
toolsDir := filepath.Join(workDir, ".tools")
|
||||
if err := os.MkdirAll(toolsDir, 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
f, err := os.CreateTemp(toolsDir, key+"-*.py")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
tmp := f.Name()
|
||||
defer os.Remove(tmp)
|
||||
if _, err := f.WriteString(code); err != nil {
|
||||
f.Close()
|
||||
return "", err
|
||||
}
|
||||
f.Close()
|
||||
|
||||
runCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
c := exec.CommandContext(runCtx, interp, tmp)
|
||||
c.Dir = workDir
|
||||
c.Env = append(os.Environ(), sessionEnv...) // 会话代理 env
|
||||
for k, v := range params { // 标量参数镜像成 TOOL_<NAME>
|
||||
if sv, ok := scalarStr(v); ok {
|
||||
c.Env = append(c.Env, "TOOL_"+strings.ToUpper(k)+"="+sv)
|
||||
}
|
||||
}
|
||||
pj, _ := json.Marshal(params)
|
||||
c.Stdin = bytes.NewReader(pj) // 参数 JSON 走 stdin
|
||||
out, err := c.CombinedOutput()
|
||||
body := string(out)
|
||||
if runCtx.Err() == context.DeadlineExceeded {
|
||||
body += "\n... [超时终止] ..."
|
||||
} else if err != nil {
|
||||
body += "\n[exit: " + err.Error() + "]"
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// ---------- http:原生请求 + 代理 ----------
|
||||
|
||||
func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, params map[string]any) (actool.Result, error) {
|
||||
var spec httpExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
method := strings.ToUpper(strings.TrimSpace(spec.Method))
|
||||
if method == "" {
|
||||
method = "GET"
|
||||
}
|
||||
rawURL := renderTemplate(spec.URL, params, identity)
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return actool.Errorf("http url 为空"), nil
|
||||
}
|
||||
var bodyReader io.Reader
|
||||
if spec.Body != "" {
|
||||
bodyReader = strings.NewReader(renderTemplate(spec.Body, params, identity))
|
||||
}
|
||||
runCtx, cancel := context.WithTimeout(ctx, timeoutOr(spec.TimeoutMs, 30000))
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(runCtx, method, rawURL, bodyReader)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
for k, v := range spec.Headers {
|
||||
req.Header.Set(k, renderTemplate(v, params, identity))
|
||||
}
|
||||
client := &http.Client{Timeout: timeoutOr(spec.TimeoutMs, 30000)}
|
||||
if tr := s.httpProxyTransport(spec); tr != nil {
|
||||
client.Transport = tr
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return actool.Errorf("请求失败: " + err.Error()), nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
out := map[string]any{"status": resp.StatusCode, "body": string(respBody)}
|
||||
b, _ := json.Marshal(out)
|
||||
return actool.Text(clipOutput(string(b), nil)), nil
|
||||
}
|
||||
|
||||
// httpProxyTransport builds a Transport for the http tool's proxy config, or nil
|
||||
// (direct). use_recording_proxy routes through the recording proxy + trusts its CA.
|
||||
func (s *Server) httpProxyTransport(spec httpExec) *http.Transport {
|
||||
proxyStr := strings.TrimSpace(spec.Proxy)
|
||||
var caFile string
|
||||
if spec.UseRecordingProxy {
|
||||
if addr := s.m.ProxyAddr(); addr != "" {
|
||||
proxyStr = "http://" + addr
|
||||
caFile = s.m.ProxyCACert()
|
||||
}
|
||||
}
|
||||
if proxyStr == "" {
|
||||
return nil
|
||||
}
|
||||
pu, err := url.Parse(proxyStr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
tr := &http.Transport{Proxy: http.ProxyURL(pu)}
|
||||
if caFile != "" {
|
||||
if pem, err := os.ReadFile(caFile); err == nil {
|
||||
pool := x509.NewCertPool()
|
||||
if pool.AppendCertsFromPEM(pem) {
|
||||
tr.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||
}
|
||||
}
|
||||
}
|
||||
return tr
|
||||
}
|
||||
|
||||
// ---------- helpers ----------
|
||||
|
||||
func identity(s string) string { return s }
|
||||
|
||||
// shellQuote single-quotes a value for safe shell interpolation.
|
||||
func shellQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" }
|
||||
|
||||
// renderTemplate replaces {name} placeholders with each param's rendered value.
|
||||
func renderTemplate(tmpl string, params map[string]any, quote func(string) string) string {
|
||||
out := tmpl
|
||||
for k, v := range params {
|
||||
out = strings.ReplaceAll(out, "{"+k+"}", quote(valToStr(v)))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// valToStr renders a param value: scalars as-is, arrays/objects as compact JSON.
|
||||
func valToStr(v any) string {
|
||||
if sv, ok := scalarStr(v); ok {
|
||||
return sv
|
||||
}
|
||||
b, _ := json.Marshal(v)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// scalarStr returns (string, true) for scalar values, ("", false) for arrays/objects.
|
||||
func scalarStr(v any) (string, bool) {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x, true
|
||||
case bool:
|
||||
return fmt.Sprintf("%t", x), true
|
||||
case float64:
|
||||
if x == float64(int64(x)) {
|
||||
return fmt.Sprintf("%d", int64(x)), true
|
||||
}
|
||||
return fmt.Sprintf("%g", x), true
|
||||
case nil:
|
||||
return "", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// clipOutput caps output to the session's MaxOutputChars (or a default).
|
||||
func clipOutput(s string, tc *actool.ToolContext) string {
|
||||
max := 6000
|
||||
if tc != nil && tc.MaxOutputChars > 0 {
|
||||
max = tc.MaxOutputChars
|
||||
}
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
return s[:max] + fmt.Sprintf("\n... [截断，共 %d 字节] ...", len(s))
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRenderTemplateCommand(t *testing.T) {
|
||||
// command: string params shell-quoted; array → JSON (also quoted).
|
||||
got := renderTemplate("nmap -p {ports} {target}", map[string]any{
|
||||
"target": "10.0.0.1; rm -rf /", // injection attempt → must be single-quoted
|
||||
"ports": []any{80.0, 443.0},
|
||||
}, shellQuote)
|
||||
if !strings.Contains(got, `'10.0.0.1; rm -rf /'`) {
|
||||
t.Fatalf("target not shell-quoted: %q", got)
|
||||
}
|
||||
if strings.Contains(got, "rm -rf /'") && !strings.Contains(got, `'10.0.0.1; rm -rf /'`) {
|
||||
t.Fatalf("possible injection leak: %q", got)
|
||||
}
|
||||
if strings.Contains(got, "{target}") || strings.Contains(got, "{ports}") {
|
||||
t.Fatalf("placeholders not replaced: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderTemplateHTTP(t *testing.T) {
|
||||
// http: identity (no shell quoting) — raw substitution into url/body.
|
||||
got := renderTemplate("https://x/submit?flag={flag}", map[string]any{"flag": "CTF{abc}"}, identity)
|
||||
if got != "https://x/submit?flag=CTF{abc}" {
|
||||
t.Fatalf("http render: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScalarStr(t *testing.T) {
|
||||
cases := []struct {
|
||||
v any
|
||||
want string
|
||||
ok bool
|
||||
}{
|
||||
{"hi", "hi", true},
|
||||
{true, "true", true},
|
||||
{float64(80), "80", true}, // integer-valued float → no decimal
|
||||
{float64(1.5), "1.5", true}, // real float
|
||||
{[]any{1, 2}, "", false}, // array → not scalar
|
||||
{map[string]any{}, "", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got, ok := scalarStr(c.v)
|
||||
if ok != c.ok || (ok && got != c.want) {
|
||||
t.Fatalf("scalarStr(%v) = (%q,%v), want (%q,%v)", c.v, got, ok, c.want, c.ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSchema(t *testing.T) {
|
||||
// empty → thin {args:string}
|
||||
m := ensureSchema(nil)
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
if _, ok := props["args"]; !ok {
|
||||
t.Fatalf("empty schema should default to args: %v", m)
|
||||
}
|
||||
// non-empty → passthrough
|
||||
raw := json.RawMessage(`{"type":"object","properties":{"target":{"type":"string"}}}`)
|
||||
m2 := ensureSchema(raw)
|
||||
p2, _ := m2["properties"].(map[string]any)
|
||||
if _, ok := p2["target"]; !ok {
|
||||
t.Fatalf("non-empty schema should pass through: %v", m2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellQuote(t *testing.T) {
|
||||
if got := shellQuote("a'b"); got != `'a'\''b'` {
|
||||
t.Fatalf("shellQuote(a'b) = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExecPython runs a real Python script end-to-end: it must read params from
|
||||
// stdin JSON and the mirrored env var, then print — verifying the whole script
|
||||
// param-passing path. Skips if no python3.
|
||||
func TestExecPython(t *testing.T) {
|
||||
interp, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("python3 unavailable")
|
||||
}
|
||||
code := `import json,sys,os
|
||||
a = json.load(sys.stdin)
|
||||
print("stdin_target=" + a["target"])
|
||||
print("env_target=" + os.environ.get("TOOL_TARGET",""))
|
||||
print("port=" + str(a["port"]))
|
||||
`
|
||||
out, err := execPython(context.Background(), interp, "test", code,
|
||||
map[string]any{"target": "example.com", "port": float64(8080)},
|
||||
t.TempDir(), nil, 10*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("execPython error: %v", err)
|
||||
}
|
||||
for _, want := range []string{"stdin_target=example.com", "env_target=example.com", "port=8080"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Fatalf("output missing %q:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunHTTPTool functionally exercises the http executor against a local server:
|
||||
// method/url/header/body templates are rendered from params, the request is sent,
|
||||
// and the response status+body are returned. No recording proxy → s.m untouched.
|
||||
func TestRunHTTPTool(t *testing.T) {
|
||||
var gotMethod, gotAuth, gotBody string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotMethod = r.Method
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
w.WriteHeader(201)
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
execRaw, _ := json.Marshal(map[string]any{
|
||||
"method": "POST",
|
||||
"url": srv.URL + "/submit?flag={flag}",
|
||||
"headers": map[string]string{"Authorization": "Bearer {token}"},
|
||||
"body": `{"flag":"{flag}"}`,
|
||||
})
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), execRaw,
|
||||
map[string]any{"flag": "CTF{x}", "token": "sekret"})
|
||||
if err != nil {
|
||||
t.Fatalf("runHTTPTool: %v", err)
|
||||
}
|
||||
if res.IsError || len(res.Content) == 0 {
|
||||
t.Fatalf("unexpected result: %+v", res)
|
||||
}
|
||||
var out struct {
|
||||
Status int `json:"status"`
|
||||
Body string `json:"body"`
|
||||
}
|
||||
_ = json.Unmarshal([]byte(res.Content[0].Text), &out)
|
||||
|
||||
if gotMethod != "POST" {
|
||||
t.Fatalf("method: %q", gotMethod)
|
||||
}
|
||||
if gotAuth != "Bearer sekret" {
|
||||
t.Fatalf("header template not rendered: %q", gotAuth)
|
||||
}
|
||||
if gotBody != `{"flag":"CTF{x}"}` {
|
||||
t.Fatalf("body template not rendered: %q", gotBody)
|
||||
}
|
||||
if out.Status != 201 || !strings.Contains(out.Body, `"ok":true`) {
|
||||
t.Fatalf("response not captured: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectPython(t *testing.T) {
|
||||
if _, err := exec.LookPath("python3"); err != nil {
|
||||
t.Skip("python3 unavailable")
|
||||
}
|
||||
if p := detectPython(); p == "" {
|
||||
t.Fatal("detectPython returned empty despite python3 on PATH")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClipOutput(t *testing.T) {
|
||||
long := strings.Repeat("x", 7000)
|
||||
got := clipOutput(long, nil)
|
||||
if len(got) >= 7000 || !strings.Contains(got, "截断") {
|
||||
t.Fatalf("clipOutput should truncate long output, got len %d", len(got))
|
||||
}
|
||||
short := "ok"
|
||||
if clipOutput(short, nil) != short {
|
||||
t.Fatal("clipOutput should pass short output through")
|
||||
}
|
||||
}
|
||||
+386
@@ -0,0 +1,386 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// DTO/serialization layer: each handler emits EXACTLY the frontend's spec shapes
|
||||
// (artex/web/src/lib/types.ts). These reshape db package structs so the
|
||||
// internal DB model never leaks over the API. The db structs and the frontend are
|
||||
// the canonical contracts; this file maps one onto the other.
|
||||
|
||||
func i64s(v int64) string { return strconv.FormatInt(v, 10) }
|
||||
|
||||
func rfc3339(t time.Time) string { return t.Format(time.RFC3339) }
|
||||
|
||||
// rawString stringifies a json.RawMessage, returning "" for empty/nil so omitempty
|
||||
// fields drop out.
|
||||
func rawString(raw json.RawMessage) string {
|
||||
if len(raw) == 0 {
|
||||
return ""
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
// ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status.
|
||||
type TaskDTO struct {
|
||||
ID string `json:"id"`
|
||||
ExplorationID int64 `json:"exploration_id"`
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
Status string `json:"status"` // created | running | paused | done | failed
|
||||
CreatedAt string `json:"created_at"`
|
||||
CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc)
|
||||
CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished
|
||||
CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished)
|
||||
LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none)
|
||||
Paused bool `json:"paused"`
|
||||
Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default
|
||||
}
|
||||
|
||||
// TokenTotalDTO is a whole-task (all agents) token aggregate.
|
||||
type TokenTotalDTO struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
|
||||
return TokenTotalDTO{
|
||||
InputTokens: u.InputTokens,
|
||||
OutputTokens: u.OutputTokens,
|
||||
CacheReadTokens: u.CacheReadTokens,
|
||||
CacheWriteTokens: u.CacheWriteTokens,
|
||||
}
|
||||
}
|
||||
|
||||
func taskDTO(t *Task, status string) TaskDTO {
|
||||
return TaskDTO{
|
||||
ID: t.ID,
|
||||
ExplorationID: t.ExpID,
|
||||
Description: t.Description,
|
||||
Goal: t.Goal,
|
||||
Status: status,
|
||||
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
|
||||
CreatedUnix: t.CreatedAt,
|
||||
CompletedAt: completedRFC(t.CompletedAt),
|
||||
CompletedUnix: t.CompletedAt,
|
||||
Paused: t.Paused,
|
||||
LLMProfileID: t.LLMProfileID,
|
||||
}
|
||||
}
|
||||
|
||||
// completedRFC renders a unix completion time as RFC3339, or "" when unset (0).
|
||||
func completedRFC(unix int64) string {
|
||||
if unix == 0 {
|
||||
return ""
|
||||
}
|
||||
return rfc3339(time.Unix(unix, 0))
|
||||
}
|
||||
|
||||
// ---- TrafficExchange (frontend "TrafficExchange") ---- ts as RFC3339.
|
||||
type TrafficExchangeDTO struct {
|
||||
ID string `json:"id"`
|
||||
TS string `json:"ts"`
|
||||
Host string `json:"host"`
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
Status int `json:"status"`
|
||||
ContentType string `json:"content_type"`
|
||||
RespLen int `json:"resp_len"`
|
||||
}
|
||||
|
||||
func trafficDTOs(ex []traffic.ExchangeMeta) []TrafficExchangeDTO {
|
||||
out := make([]TrafficExchangeDTO, 0, len(ex))
|
||||
for _, e := range ex {
|
||||
out = append(out, TrafficExchangeDTO{
|
||||
ID: e.ID,
|
||||
TS: rfc3339(time.Unix(e.TS, 0)),
|
||||
Host: e.Host,
|
||||
Method: e.Method,
|
||||
URL: e.URL,
|
||||
Status: e.Status,
|
||||
ContentType: e.ContentType,
|
||||
RespLen: e.RespLen,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- TaskNode (frontend "TaskNode") — frontier, intents, exploration graph nodes ----
|
||||
|
||||
type TaskNodeDTO struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // db Node.Kind
|
||||
Payload string `json:"payload,omitempty"`
|
||||
Priority int `json:"priority"`
|
||||
State string `json:"state"`
|
||||
Origin string `json:"origin"`
|
||||
TS string `json:"ts"`
|
||||
}
|
||||
|
||||
func taskNodeDTO(n *db.Node) TaskNodeDTO {
|
||||
return TaskNodeDTO{
|
||||
ID: i64s(n.ID),
|
||||
Type: n.Kind,
|
||||
Payload: rawString(n.Payload),
|
||||
Priority: n.Priority,
|
||||
State: n.State,
|
||||
Origin: n.Origin,
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
}
|
||||
|
||||
func taskNodeDTOs(in []*db.Node) []TaskNodeDTO {
|
||||
out := make([]TaskNodeDTO, 0, len(in))
|
||||
for _, n := range in {
|
||||
out = append(out, taskNodeDTO(n))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Edge (frontend "Edge") — exploration edges and asset edges ----
|
||||
|
||||
type EdgeDTO struct {
|
||||
Src string `json:"src"`
|
||||
Dst string `json:"dst"`
|
||||
Rel string `json:"rel"`
|
||||
}
|
||||
|
||||
func edgeDTO(e db.Edge) EdgeDTO {
|
||||
return EdgeDTO{Src: i64s(e.From), Dst: i64s(e.To), Rel: e.Rel}
|
||||
}
|
||||
|
||||
func edgeDTOs(in []db.Edge) []EdgeDTO {
|
||||
out := make([]EdgeDTO, 0, len(in))
|
||||
for _, e := range in {
|
||||
out = append(out, edgeDTO(e))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
ID string `json:"id"`
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Severity string `json:"severity"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence string `json:"evidence"`
|
||||
IntentID string `json:"intent_id,omitempty"`
|
||||
ParamID string `json:"param_id,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
TaskDescription string `json:"task_description,omitempty"`
|
||||
TS string `json:"ts"`
|
||||
}
|
||||
|
||||
// findingPayload mirrors the JSON written by the worker's report_finding tool
|
||||
// (agent/tools.go addFinding): {vulnclass, severity, summary, evidence:{by,poc}}.
|
||||
type findingPayload struct {
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Severity string `json:"severity"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence json.RawMessage `json:"evidence"`
|
||||
}
|
||||
|
||||
func findingDTO(n *db.Node) FindingDTO {
|
||||
var p findingPayload
|
||||
_ = json.Unmarshal(n.Payload, &p)
|
||||
return FindingDTO{
|
||||
ID: i64s(n.ID),
|
||||
VulnClass: p.VulnClass,
|
||||
Severity: p.Severity,
|
||||
Summary: p.Summary,
|
||||
Evidence: rawString(p.Evidence),
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
}
|
||||
|
||||
// findingDTOsForTask converts a task's finding nodes to DTOs, stamping each with
|
||||
// the owning task's id/description so the global 发现 page can group across tasks.
|
||||
func findingDTOsForTask(t *Task, in []*db.Node) []FindingDTO {
|
||||
out := make([]FindingDTO, 0, len(in))
|
||||
for _, n := range in {
|
||||
d := findingDTO(n)
|
||||
d.TaskID = t.ID
|
||||
d.TaskDescription = t.Description
|
||||
out = append(out, d)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// findingFromDB converts a standalone DBFinding row to a FindingDTO. task_id and
|
||||
// task_description are empty when the originating task has been deleted (NULL).
|
||||
func findingFromDB(f *db.DBFinding) FindingDTO {
|
||||
d := FindingDTO{
|
||||
ID: i64s(f.ID),
|
||||
VulnClass: f.VulnClass,
|
||||
Severity: f.Severity,
|
||||
Summary: f.Summary,
|
||||
Evidence: f.Evidence,
|
||||
TS: rfc3339(f.CreatedAt),
|
||||
}
|
||||
if f.TaskID != nil {
|
||||
d.TaskID = i64s(*f.TaskID)
|
||||
d.TaskDescription = f.TaskDescription
|
||||
}
|
||||
if f.NodeID != nil {
|
||||
d.IntentID = "" // node_id is the finding node, not the intent; keep IntentID empty
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// ---- Activity (frontend "Activity") ----
|
||||
|
||||
type ActivityDTO struct {
|
||||
Seq int64 `json:"seq"` // db Activity.ID
|
||||
IntentID string `json:"intent_id,omitempty"` // db NodeID
|
||||
Worker string `json:"worker"`
|
||||
TS string `json:"ts"` // db CreatedAt
|
||||
Kind string `json:"kind"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
IsError bool `json:"is_error"`
|
||||
Summary string `json:"summary"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
// token usage (set only on kind='result'); used for per-session token totals.
|
||||
InputTokens *int `json:"input_tokens,omitempty"`
|
||||
OutputTokens *int `json:"output_tokens,omitempty"`
|
||||
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
|
||||
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
|
||||
}
|
||||
|
||||
func activityDTO(a db.Activity) ActivityDTO {
|
||||
intent := ""
|
||||
if a.NodeID != nil {
|
||||
intent = i64s(*a.NodeID)
|
||||
}
|
||||
return ActivityDTO{
|
||||
Seq: a.ID,
|
||||
IntentID: intent,
|
||||
Worker: a.Worker,
|
||||
TS: rfc3339(a.CreatedAt),
|
||||
Kind: a.Kind,
|
||||
Tool: a.Tool,
|
||||
ToolUseID: a.ToolUseID,
|
||||
IsError: a.IsError,
|
||||
Summary: a.Summary,
|
||||
Detail: a.Detail, // list endpoint leaves this empty (lazy)
|
||||
InputTokens: a.InputTokens,
|
||||
OutputTokens: a.OutputTokens,
|
||||
CacheReadTokens: a.CacheReadTokens,
|
||||
CacheWriteTokens: a.CacheWriteTokens,
|
||||
}
|
||||
}
|
||||
|
||||
func activityDTOs(in []db.Activity) []ActivityDTO {
|
||||
out := make([]ActivityDTO, 0, len(in))
|
||||
for _, a := range in {
|
||||
out = append(out, activityDTO(a))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Agent (frontend "Agent") — string id ----
|
||||
|
||||
type AgentDTO struct {
|
||||
ID string `json:"id"`
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Role string `json:"role"`
|
||||
Builtin bool `json:"builtin"`
|
||||
Enabled bool `json:"enabled"`
|
||||
MaxTurns int `json:"max_turns"`
|
||||
RunSecs int `json:"run_seconds"`
|
||||
WebSearch bool `json:"web_search"`
|
||||
InteractiveShell bool `json:"interactive_shell"`
|
||||
// binding counts (populated only by the list endpoint) — shown on agent cards.
|
||||
McpCount int `json:"mcp_count"`
|
||||
SkillCount int `json:"skill_count"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
}
|
||||
|
||||
func agentDTO(a *db.Agent) AgentDTO {
|
||||
return AgentDTO{
|
||||
ID: i64s(a.ID),
|
||||
Key: a.Key,
|
||||
Name: a.Name,
|
||||
Description: a.Description,
|
||||
Role: a.Role,
|
||||
Builtin: a.Builtin,
|
||||
Enabled: a.Enabled,
|
||||
MaxTurns: a.MaxTurns,
|
||||
RunSecs: a.RunSecs,
|
||||
WebSearch: a.WebSearch,
|
||||
InteractiveShell: a.InteractiveShell,
|
||||
}
|
||||
}
|
||||
|
||||
func agentDTOs(in []*db.Agent) []AgentDTO {
|
||||
out := make([]AgentDTO, 0, len(in))
|
||||
for _, a := range in {
|
||||
out = append(out, agentDTO(a))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- LLMProfile (frontend "LLMProfile") — string id ----
|
||||
|
||||
type LLMProfileDTO struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Format string `json:"format"`
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
Proxy string `json:"proxy,omitempty"`
|
||||
Model string `json:"model"`
|
||||
APIKeyHint string `json:"api_key_hint,omitempty"`
|
||||
RatePerSecond float64 `json:"rate_per_second"`
|
||||
RatePerMinute float64 `json:"rate_per_minute"`
|
||||
ContextWindowK int `json:"context_window_k"`
|
||||
ReasoningEffort string `json:"reasoning_effort"`
|
||||
IsDefault bool `json:"is_default"`
|
||||
}
|
||||
|
||||
func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO {
|
||||
return LLMProfileDTO{
|
||||
ID: i64s(p.ID),
|
||||
Name: p.Name,
|
||||
Format: p.Format,
|
||||
BaseURL: p.BaseURL,
|
||||
Proxy: p.Proxy,
|
||||
Model: p.Model,
|
||||
APIKeyHint: p.APIKeyHint,
|
||||
RatePerSecond: p.RatePerSecond,
|
||||
RatePerMinute: p.RatePerMinute,
|
||||
ContextWindowK: p.ContextWindowK,
|
||||
ReasoningEffort: p.ReasoningEffort,
|
||||
IsDefault: p.IsDefault,
|
||||
}
|
||||
}
|
||||
|
||||
func llmProfileDTOs(in []*db.LLMProfile) []LLMProfileDTO {
|
||||
out := make([]LLMProfileDTO, 0, len(in))
|
||||
for _, p := range in {
|
||||
out = append(out, llmProfileDTO(p))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// idStrings formats a slice of int64 ids as strings (for visibility agent ids).
|
||||
func idStrings(in []int64) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
for _, v := range in {
|
||||
out = append(out, i64s(v))
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,595 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
// model_error（provider/API 故障：LLM 层瞬时重试耗尽，或流已开始后中途断流）
|
||||
// 收场的 work 不是「试过没做完」也不是真失败，而是外部抖动。默认把它当永久
|
||||
// blocked 会白丢一条意图，所以这里对该终态额外重跑几次，每次之间退避一下，给
|
||||
// provider 恢复的时间；重试期间若被暂停/终止/取消则立即让位给对应分支处理。
|
||||
const (
|
||||
modelErrorRetries = 2 // model_error 收场后额外重试的次数
|
||||
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
|
||||
)
|
||||
|
||||
// Engine drives the event-driven exploration loop with real LLM agents
|
||||
// (docs §4.3/§4.4): on asset/exploration-graph change (debounced) it wakes the
|
||||
// planner, which reads the route, queries assets, judges goals and emits intents;
|
||||
// N concurrent work agents claim intents and execute them. There is no
|
||||
// simulation mode — an LLM provider is required. The planner/worker can be
|
||||
// (re)installed at runtime (LLM configured from the UI); the loops always run
|
||||
// but idle until an LLM is set.
|
||||
type Engine struct {
|
||||
m *Manager
|
||||
debounce time.Duration
|
||||
|
||||
bc *Broadcaster // live activity pub/sub (SSE)
|
||||
|
||||
mu sync.RWMutex
|
||||
planner *agent.Planner
|
||||
worker *agent.Worker
|
||||
started sync.Map // taskID -> bool, so Run is idempotent per task
|
||||
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
||||
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
|
||||
|
||||
// per-task execution context: each planner.Plan / worker.Execute runs under it,
|
||||
// so pausing can CANCEL an in-flight run (not just skip the next one). Recreated
|
||||
// on resume since cancelling is one-shot.
|
||||
execMu sync.Mutex
|
||||
execCancel map[string]context.CancelFunc
|
||||
execCtx map[string]context.Context
|
||||
|
||||
// per-work (per running intent) cancel, so the planner's kill_work can stop a
|
||||
// single worker without pausing the whole task. Keyed by globally-unique intent id.
|
||||
workMu sync.Mutex
|
||||
workCancel map[int64]context.CancelFunc
|
||||
|
||||
// steerBox queues planner course-corrections for a running work (keyed by intent
|
||||
// id). The worker's PreToolUse hook drains it before its next tool call and hands
|
||||
// the message to the model (blocking that call) so it re-plans — no kill needed.
|
||||
steerMu sync.Mutex
|
||||
steerBox map[int64][]string
|
||||
|
||||
plannerRound sync.Map // taskID -> int, planner round counter (for UI round separators)
|
||||
|
||||
// 任务级超时(见 docs/任务级超时与收尾设计.md):
|
||||
settling sync.Map // taskID -> bool, 任务已进入收尾时序(停止派/领新意图)
|
||||
deadline sync.Map // taskID -> int64 unix, 绝对截止时刻(首次运行时盖章;0/缺省=不限)
|
||||
stamped sync.Map // taskID -> bool, first_run_at 是否已盖章(本进程内只盖一次)
|
||||
inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain)
|
||||
coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重)
|
||||
|
||||
// resolve, if set, returns a task's dedicated planner/worker when it pins a
|
||||
// specific LLM profile (nil,nil = use the global active pair). Set once at startup.
|
||||
resolve func(t *Task) (*agent.Planner, *agent.Worker)
|
||||
}
|
||||
|
||||
// nextPlannerRound returns the next planner round number for a task (1-based).
|
||||
func (e *Engine) nextPlannerRound(taskID string) int {
|
||||
v, _ := e.plannerRound.LoadOrStore(taskID, 0)
|
||||
n := v.(int) + 1
|
||||
e.plannerRound.Store(taskID, n)
|
||||
return n
|
||||
}
|
||||
|
||||
// Pause stops a task: marks it paused AND cancels any in-flight planner/worker run
|
||||
// for it (a long worker.Execute would otherwise keep going until it finishes).
|
||||
func (e *Engine) Pause(taskID string) {
|
||||
e.paused.Store(taskID, true)
|
||||
e.cancelExec(taskID)
|
||||
}
|
||||
|
||||
// cancelExec cancels a task's current per-task exec context (any in-flight
|
||||
// planner.Plan / worker.Execute), if present. Shared by Pause and the settle
|
||||
// sequence's hard-drain backstop.
|
||||
func (e *Engine) cancelExec(taskID string) {
|
||||
e.execMu.Lock()
|
||||
if cancel := e.execCancel[taskID]; cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
e.execMu.Unlock()
|
||||
}
|
||||
|
||||
// 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) {
|
||||
e.paused.Delete(t.ID)
|
||||
t.Notify()
|
||||
}
|
||||
|
||||
// execContextFor returns a live per-task context derived from parent, recreating
|
||||
// it if a prior pause cancelled it.
|
||||
func (e *Engine) execContextFor(parent context.Context, taskID string) context.Context {
|
||||
e.execMu.Lock()
|
||||
defer e.execMu.Unlock()
|
||||
if e.IsPaused(taskID) {
|
||||
// never hand out a live context while paused (guards the claim→Execute race)
|
||||
c, cancel := context.WithCancel(parent)
|
||||
cancel()
|
||||
return c
|
||||
}
|
||||
if c := e.execCtx[taskID]; c != nil && c.Err() == nil {
|
||||
return c
|
||||
}
|
||||
c, cancel := context.WithCancel(parent)
|
||||
e.execCtx[taskID] = c
|
||||
e.execCancel[taskID] = cancel
|
||||
return c
|
||||
}
|
||||
|
||||
// IsPaused reports whether a task is user-paused.
|
||||
func (e *Engine) IsPaused(taskID string) bool {
|
||||
v, ok := e.paused.Load(taskID)
|
||||
return ok && v.(bool)
|
||||
}
|
||||
|
||||
// Started reports whether the engine loops are running for a task.
|
||||
func (e *Engine) Started(taskID string) bool {
|
||||
_, ok := e.started.Load(taskID)
|
||||
return ok
|
||||
}
|
||||
|
||||
// LastActivity returns the unix time of the last planner/worker activity for a
|
||||
// task (0 if none yet).
|
||||
func (e *Engine) LastActivity(taskID string) int64 {
|
||||
if v, ok := e.lastAct.Load(taskID); ok {
|
||||
return v.(int64)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) }
|
||||
|
||||
func NewEngine(m *Manager) *Engine {
|
||||
return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(),
|
||||
execCancel: map[string]context.CancelFunc{}, execCtx: map[string]context.Context{},
|
||||
workCancel: map[int64]context.CancelFunc{}, steerBox: map[int64][]string{}}
|
||||
}
|
||||
|
||||
// registerWork records the cancel for the work currently running intentID.
|
||||
func (e *Engine) registerWork(intentID int64, cancel context.CancelFunc) {
|
||||
e.workMu.Lock()
|
||||
e.workCancel[intentID] = cancel
|
||||
e.workMu.Unlock()
|
||||
}
|
||||
|
||||
// unregisterWork removes and releases the work's cancel (called when Execute returns).
|
||||
func (e *Engine) unregisterWork(intentID int64) {
|
||||
e.workMu.Lock()
|
||||
if c := e.workCancel[intentID]; c != nil {
|
||||
delete(e.workCancel, intentID)
|
||||
c() // release context resources (no-op if already cancelled)
|
||||
}
|
||||
e.workMu.Unlock()
|
||||
e.steerMu.Lock()
|
||||
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
|
||||
e.steerMu.Unlock()
|
||||
}
|
||||
|
||||
// SteerWork queues a mid-run course-correction for the work running intentID (the
|
||||
// planner's steer_work tool). The worker delivers it before its next tool call and
|
||||
// re-plans — no kill. Errors if no work is currently running that intent.
|
||||
func (e *Engine) SteerWork(intentID int64, msg string) error {
|
||||
if strings.TrimSpace(msg) == "" {
|
||||
return fmt.Errorf("纠偏消息不能为空")
|
||||
}
|
||||
e.workMu.Lock()
|
||||
running := e.workCancel[intentID] != nil
|
||||
e.workMu.Unlock()
|
||||
if !running {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID)
|
||||
}
|
||||
e.steerMu.Lock()
|
||||
e.steerBox[intentID] = append(e.steerBox[intentID], msg)
|
||||
e.steerMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// drainSteer pops the oldest queued steering message for intentID (FIFO), if any.
|
||||
func (e *Engine) drainSteer(intentID int64) (string, bool) {
|
||||
e.steerMu.Lock()
|
||||
defer e.steerMu.Unlock()
|
||||
q := e.steerBox[intentID]
|
||||
if len(q) == 0 {
|
||||
return "", false
|
||||
}
|
||||
msg := q[0]
|
||||
if len(q) == 1 {
|
||||
delete(e.steerBox, intentID)
|
||||
} else {
|
||||
e.steerBox[intentID] = q[1:]
|
||||
}
|
||||
return msg, true
|
||||
}
|
||||
|
||||
// steerHooks wraps the guard's hook runner so the planner can steer a running work:
|
||||
// before each tool call it drains a queued course-correction (if any) and blocks the
|
||||
// call, handing the message back to the model — which re-plans its next step instead
|
||||
// of running the tool. No queued message → the guard behaves exactly as before.
|
||||
type steerHooks struct {
|
||||
inner harness.HookRunner
|
||||
drain func() (string, bool)
|
||||
}
|
||||
|
||||
func (h steerHooks) PreToolUse(ctx context.Context, name string, input []byte) (bool, string, []byte) {
|
||||
if msg, ok := h.drain(); ok {
|
||||
return true, "【规划者实时纠偏】" + msg +
|
||||
"\n（这是规划者对本意图的即时指令；本次工具调用未执行，请据此调整下一步。若与你当前打算冲突，以此为准。）", nil
|
||||
}
|
||||
if h.inner != nil {
|
||||
return h.inner.PreToolUse(ctx, name, input)
|
||||
}
|
||||
return false, "", nil
|
||||
}
|
||||
|
||||
func (h steerHooks) PostToolUse(ctx context.Context, name string, input, result []byte, isErr bool) {
|
||||
if h.inner != nil {
|
||||
h.inner.PostToolUse(ctx, name, input, result, isErr)
|
||||
}
|
||||
}
|
||||
|
||||
func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []string, string) {
|
||||
if h.inner != nil {
|
||||
return h.inner.Stop(ctx, messages)
|
||||
}
|
||||
return false, nil, ""
|
||||
}
|
||||
|
||||
// KillWork cancels the in-flight work running intentID (planner's kill_work tool).
|
||||
// The work's agent-core session honors ctx cancellation and aborts promptly.
|
||||
func (e *Engine) KillWork(intentID int64) error {
|
||||
e.workMu.Lock()
|
||||
c := e.workCancel[intentID]
|
||||
e.workMu.Unlock()
|
||||
if c == nil {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID)
|
||||
}
|
||||
c()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcaster exposes the engine's live activity pub/sub (used by the SSE handler).
|
||||
func (e *Engine) Broadcaster() *Broadcaster { return e.bc }
|
||||
|
||||
// emitActivity persists one captured step AND fans it out to live subscribers,
|
||||
// from a single point so storage and the SSE stream never diverge.
|
||||
func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
||||
id, err := e.appendActivity(t, r)
|
||||
if err != nil {
|
||||
// NO LONGER SILENT: dropping a record breaks command↔result pairing in the
|
||||
// trace — a tool_use whose tool_result was lost shows as "执行中" forever.
|
||||
log.Printf("[activity] task %s 重试后仍失败、丢弃该记录 (kind=%s tool=%s tuid=%s): %v",
|
||||
t.ID, r.Kind, r.Tool, r.ToolUseID, err)
|
||||
e.touch(t.ID)
|
||||
return
|
||||
}
|
||||
r.ID = id
|
||||
if r.CreatedAt.IsZero() {
|
||||
r.CreatedAt = time.Now()
|
||||
}
|
||||
e.bc.Publish(t.ID, r)
|
||||
e.touch(t.ID)
|
||||
}
|
||||
|
||||
// appendActivity persists one activity row, retrying briefly on write failure.
|
||||
// Concurrent planner + worker inserts into the same exploration's activity log
|
||||
// occasionally fail; a couple of quick retries recover most. Crucially, every
|
||||
// failure is now LOGGED (it used to be swallowed by an `if err == nil`), so the
|
||||
// underlying DB error is finally visible for diagnosis.
|
||||
func (e *Engine) appendActivity(t *Task, r db.Activity) (int64, error) {
|
||||
var id int64
|
||||
var err error
|
||||
for attempt := 1; attempt <= 3; attempt++ {
|
||||
if id, err = t.Store.AppendActivity(r); err == nil {
|
||||
if attempt > 1 {
|
||||
log.Printf("[activity] task %s 写入第 %d 次重试成功 (worker=%s kind=%s tool=%s)",
|
||||
t.ID, attempt, r.Worker, r.Kind, r.Tool)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
log.Printf("[activity] task %s 写入失败 (第 %d/3 次, worker=%s kind=%s tool=%s): %v",
|
||||
t.ID, attempt, r.Worker, r.Kind, r.Tool, err)
|
||||
time.Sleep(time.Duration(attempt) * 25 * time.Millisecond)
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// UseLLM installs (or replaces) the LLM planner + work agent at runtime.
|
||||
func (e *Engine) UseLLM(p *agent.Planner, w *agent.Worker) {
|
||||
e.mu.Lock()
|
||||
e.planner, e.worker = p, w
|
||||
e.mu.Unlock()
|
||||
}
|
||||
|
||||
func (e *Engine) snapshot() (*agent.Planner, *agent.Worker) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.planner, e.worker
|
||||
}
|
||||
|
||||
// SetAgentResolver installs a per-task planner/worker resolver (wired by the server).
|
||||
// Called once at startup before any task loop runs, so no lock is needed on reads.
|
||||
func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) {
|
||||
e.resolve = fn
|
||||
}
|
||||
|
||||
// snapshotFor returns the planner/worker a task should run on: the task's dedicated
|
||||
// pair when it pins a specific LLM profile (via the resolver), else the global active
|
||||
// pair. The resolver returning nil means "no override" (fall back to global).
|
||||
func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) {
|
||||
if e.resolve != nil {
|
||||
if p, w := e.resolve(t); p != nil && w != nil {
|
||||
return p, w
|
||||
}
|
||||
}
|
||||
return e.snapshot()
|
||||
}
|
||||
|
||||
// Ready reports whether an LLM provider is configured.
|
||||
func (e *Engine) Ready() bool {
|
||||
p, w := e.snapshot()
|
||||
return p != nil && w != nil
|
||||
}
|
||||
|
||||
// Mode reports "llm" when ready, else "idle".
|
||||
func (e *Engine) Mode() string {
|
||||
if e.Ready() {
|
||||
return "llm"
|
||||
}
|
||||
return "idle"
|
||||
}
|
||||
|
||||
// Run starts the planner loop + N worker loops for a task. The loops always run
|
||||
// but no-op until an LLM is configured (so a task created while idle picks up
|
||||
// automatically once LLM is set from the UI).
|
||||
func (e *Engine) Run(ctx context.Context, t *Task) {
|
||||
if _, loaded := e.started.LoadOrStore(t.ID, true); loaded {
|
||||
t.Notify() // already running — just nudge a planning round
|
||||
return
|
||||
}
|
||||
e.touch(t.ID)
|
||||
go e.plannerLoop(ctx, t)
|
||||
workers := e.m.Workers()
|
||||
for i := 0; i < workers; i++ {
|
||||
go e.workerLoop(ctx, t, fmt.Sprintf("work#%d", i+1))
|
||||
}
|
||||
e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重)
|
||||
t.Notify() // kick the first planning round (acted on once LLM is ready)
|
||||
}
|
||||
|
||||
func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.notify:
|
||||
// debounce: coalesce a burst of changes into one planning round
|
||||
timer := time.NewTimer(e.debounce)
|
||||
drain:
|
||||
for {
|
||||
select {
|
||||
case <-t.notify:
|
||||
case <-timer.C:
|
||||
break drain
|
||||
}
|
||||
}
|
||||
planner, _ := e.snapshotFor(t)
|
||||
if planner == nil {
|
||||
continue // idle until LLM configured
|
||||
}
|
||||
if e.IsPaused(t.ID) {
|
||||
continue // user-paused: don't plan
|
||||
}
|
||||
// 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) {
|
||||
continue
|
||||
}
|
||||
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
|
||||
// 触发常规规划轮;终局那一轮由协调器(settleTask)直接驱动,不走这里。
|
||||
if e.isSettling(t.ID) {
|
||||
continue
|
||||
}
|
||||
e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
log.Printf("[planner] task %s 规划中…", t.ID)
|
||||
// round marker: each Plan() is one planner round; emit a boundary so the
|
||||
// UI can separate rounds in the transcript (kind='round').
|
||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: fmt.Sprintf("第 %d 轮规划", e.nextPlannerRound(t.ID))})
|
||||
// which intents' completion triggered this round (may be several — debounce
|
||||
// coalesces a burst; empty for non-completion wakes like a new hint).
|
||||
doneIntents := t.drainDone()
|
||||
e.incInflight(t.ID)
|
||||
taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
met, reason, err := planner.Plan(ectx, taskIDInt, e.m.assets, t.Store, t.Goal, doneIntents, emit)
|
||||
e.decInflight(t.ID)
|
||||
switch {
|
||||
case err != nil && ectx.Err() == nil:
|
||||
log.Printf("[planner] task %s 规划出错: %v", t.ID, err)
|
||||
case met:
|
||||
log.Printf("[planner] task %s 判定目标达成: %s", t.ID, reason)
|
||||
// 所有目标达成 → 持久化任务状态为 done（前端 DTO 会优先展示该终态）。
|
||||
if err := e.m.SetTaskStatus(t.ID, "done"); err != nil {
|
||||
log.Printf("[planner] task %s 标记完成落库失败: %v", t.ID, err)
|
||||
}
|
||||
// 任务已判完成 → 立刻取消在跑的 worker：它们手头的意图跑出来也没意义了。
|
||||
// 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支
|
||||
// 归为 stopped(而非 blocked)。
|
||||
e.cancelExec(t.ID)
|
||||
default:
|
||||
log.Printf("[planner] task %s 规划完成", t.ID)
|
||||
}
|
||||
e.touch(t.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
_, worker := e.snapshotFor(t)
|
||||
if worker == nil {
|
||||
if sleepCtx(ctx, 1500*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if e.IsPaused(t.ID) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // user-paused: don't claim/execute intents
|
||||
}
|
||||
if e.isSettling(t.ID) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // 任务超时收尾中:不再领新意图(在跑的自行收尾,协调器等其 drain)
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier
|
||||
}
|
||||
intent := e.claimNext(t, name)
|
||||
if intent == nil {
|
||||
if sleepCtx(ctx, 800*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID)
|
||||
e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
// per-work child context so the planner's kill_work can stop just this work.
|
||||
workCtx, workCancel := context.WithCancel(ectx)
|
||||
e.registerWork(intent.ID, workCancel)
|
||||
// wrap the guard hooks so steer_work can inject a mid-run course-correction
|
||||
// for THIS intent (drained before the worker's next tool call).
|
||||
iid := intent.ID
|
||||
taskEmit := func(a db.Activity) {
|
||||
nid := iid
|
||||
a.NodeID, a.Worker = &nid, name
|
||||
emit(a)
|
||||
}
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit)
|
||||
hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }}
|
||||
wTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.incInflight(t.ID) // 计入在跑,供收尾时序 drain 等待
|
||||
reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich)
|
||||
// model_error 收场 → 额外重跑几次（退避后再试）。仅在意图仍属本 work、任务
|
||||
// 未暂停/未终止/未取消【且未进入收尾】时重试；否则让位给对应分支处理(收尾期不
|
||||
// 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。
|
||||
for attempt := 1; attempt <= modelErrorRetries &&
|
||||
reason == harness.ReasonModelError &&
|
||||
workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ {
|
||||
log.Printf("[worker %s] task %s 意图 #%d model_error 收场，%v 后重试 (%d/%d)",
|
||||
name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries)
|
||||
if sleepCtx(workCtx, modelErrorRetryBackoff) {
|
||||
break // 退避期间被取消（终止/暂停）→ 交给下方分支处理
|
||||
}
|
||||
reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich)
|
||||
}
|
||||
e.decInflight(t.ID)
|
||||
// CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's
|
||||
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
||||
// 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.
|
||||
killed := workCtx.Err() != nil && ectx.Err() == nil
|
||||
e.unregisterWork(intent.ID)
|
||||
// if a pause cancelled this run mid-flight, return the intent to the frontier
|
||||
// so it is re-claimed (and re-run from scratch) on resume — not marked done.
|
||||
if ectx.Err() != nil && e.IsPaused(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "open")
|
||||
continue
|
||||
}
|
||||
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
|
||||
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
|
||||
if ectx.Err() != nil && e.isSettling(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "exhausted")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted)，写回 %d 条", name, t.ID, intent.ID, wrote)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
|
||||
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
|
||||
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
|
||||
if killed {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
t.Notify()
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
log.Printf("[worker %s] intent %d: %v", name, intent.ID, err)
|
||||
}
|
||||
// terminal分流：撞步数上限 ≠ 完成。max_turns→exhausted（规划者据此知道这个方向
|
||||
// 试过但没真正做完、需换角度，而非当成已覆盖永久跳过）；出错→blocked；正常→done。
|
||||
state := "done"
|
||||
switch {
|
||||
case err != nil:
|
||||
state = "blocked"
|
||||
case reason == harness.ReasonMaxTurns:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 撞步数上限(exhausted)，本次写回事实 %d 条", name, intent.ID, wrote)
|
||||
case reason == harness.ReasonTimeout:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 运行超时(exhausted)，收尾后写回事实 %d 条", name, intent.ID, wrote)
|
||||
}
|
||||
_ = t.Store.SetIntentState(intent.ID, state)
|
||||
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %d 条事实)", name, t.ID, intent.ID, state, wrote)
|
||||
e.touch(t.ID)
|
||||
t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id)
|
||||
}
|
||||
}
|
||||
|
||||
func sleepCtx(ctx context.Context, d time.Duration) (done bool) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return true
|
||||
case <-time.After(d):
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) claimNext(t *Task, name string) *db.Node {
|
||||
fr, _ := t.Store.Frontier(20)
|
||||
for _, in := range fr {
|
||||
if ok, _ := t.Store.ClaimIntent(in.ID, name); ok {
|
||||
return in
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 任务级超时协调器(见 docs/任务级超时与收尾设计.md §4/§8.5)。
|
||||
// 绝对墙钟:每个带 timeout 的任务一个定时 goroutine,到点驱动有序收尾时序:
|
||||
// ① settling → ② worker 停领新意图 / ③ planner 丢弃普通 notify
|
||||
// ④ 等在跑 worker drain(受 grace) → ⑤ 终局一轮 planner 判定 → ⑥ 定终态(带守卫)
|
||||
|
||||
const (
|
||||
settleDrainGrace = 90 * time.Second // 等在跑 worker 优雅收尾的上限;超过则硬 cancel
|
||||
deadlinePollInterval = 2 * time.Second // deadline 未盖章/LLM 未就绪时的轮询间隔
|
||||
deadlineMaxSleep = 30 * time.Second // 单次最长睡眠(便于周期复查终态)
|
||||
)
|
||||
|
||||
// ---------- settling 状态 ----------
|
||||
|
||||
func (e *Engine) isSettling(taskID string) bool {
|
||||
v, _ := e.settling.Load(taskID)
|
||||
b, _ := v.(bool)
|
||||
return b
|
||||
}
|
||||
|
||||
// markSettling flips settling on; returns true only for the first caller.
|
||||
func (e *Engine) markSettling(taskID string) bool {
|
||||
_, loaded := e.settling.LoadOrStore(taskID, true)
|
||||
return !loaded
|
||||
}
|
||||
|
||||
// ---------- 在跑计数(worker.Execute + planner.Plan),用于 drain ----------
|
||||
|
||||
func (e *Engine) inflightCounter(taskID string) *int64 {
|
||||
v, _ := e.inflight.LoadOrStore(taskID, new(int64))
|
||||
return v.(*int64)
|
||||
}
|
||||
func (e *Engine) incInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), 1) }
|
||||
func (e *Engine) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) }
|
||||
func (e *Engine) inflightCount(taskID string) int64 {
|
||||
return atomic.LoadInt64(e.inflightCounter(taskID))
|
||||
}
|
||||
|
||||
// ---------- deadline ----------
|
||||
|
||||
// taskDeadline returns the task's absolute deadline (unix). Prefers the in-process
|
||||
// map (stamped this session); falls back to the DB-loaded value (restart), seeding
|
||||
// the map. 0 = no timeout / not yet stamped.
|
||||
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
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// 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) {
|
||||
if t.TimeoutSeconds <= 0 {
|
||||
return
|
||||
}
|
||||
if _, loaded := e.stamped.LoadOrStore(t.ID, true); loaded {
|
||||
return
|
||||
}
|
||||
dl, err := e.m.StampTaskFirstRun(t.ID)
|
||||
if err != nil {
|
||||
log.Printf("[deadline] task %s 盖章 first_run 失败: %v", t.ID, err)
|
||||
e.stamped.Delete(t.ID) // 允许下次重试
|
||||
return
|
||||
}
|
||||
if dl > 0 {
|
||||
e.deadline.Store(t.ID, dl)
|
||||
log.Printf("[deadline] task %s 首次运行,截止于 %s", t.ID, time.Unix(dl, 0).Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
}
|
||||
|
||||
// clockCtx layers the task's TaskClock (absolute deadline) onto a run's context so
|
||||
// worker/planner can clamp their wall-clock budget and pick per-run vs task-timeout
|
||||
// wrap-up words. final marks the coordinator-driven terminal planner round.
|
||||
func (e *Engine) clockCtx(base context.Context, t *Task, final bool) context.Context {
|
||||
dl := e.taskDeadline(t)
|
||||
if dl <= 0 && !final {
|
||||
return base // no timeout → unchanged behavior
|
||||
}
|
||||
return agent.WithTaskClock(base, agent.TaskClock{DeadlineUnix: dl, Final: final})
|
||||
}
|
||||
|
||||
// ---------- 协调器 ----------
|
||||
|
||||
// startDeadlineCoordinator launches the per-task deadline timer once (idempotent).
|
||||
// Called from Run() and from the restart reload path, so non-active timeout tasks
|
||||
// still get settled after their deadline even without live planner/worker loops.
|
||||
func (e *Engine) startDeadlineCoordinator(ctx context.Context, t *Task) {
|
||||
if t == nil || t.TimeoutSeconds <= 0 {
|
||||
return
|
||||
}
|
||||
if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded {
|
||||
return
|
||||
}
|
||||
go e.deadlineCoordinator(ctx, t)
|
||||
}
|
||||
|
||||
// deadlineCoordinator waits until the task's absolute deadline, then runs the settle
|
||||
// sequence. Absolute wall-clock: it keeps counting through pauses.
|
||||
func (e *Engine) deadlineCoordinator(ctx context.Context, t *Task) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
return // already finished (goals met / failed) — nothing to time out
|
||||
}
|
||||
dl := e.taskDeadline(t)
|
||||
if dl <= 0 {
|
||||
if sleepCtx(ctx, deadlinePollInterval) { // not yet stamped (task hasn't really run)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if remaining := time.Until(time.Unix(dl, 0)); remaining > 0 {
|
||||
nap := remaining
|
||||
if nap > deadlineMaxSleep {
|
||||
nap = deadlineMaxSleep
|
||||
}
|
||||
if sleepCtx(ctx, nap) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
e.settleTask(ctx, t)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// settleTask runs the ordered settle sequence once (§4 steps ①–⑥).
|
||||
func (e *Engine) settleTask(ctx context.Context, t *Task) {
|
||||
if !e.markSettling(t.ID) {
|
||||
return
|
||||
}
|
||||
log.Printf("[deadline] task %s 到达超时上限,进入收尾时序", t.ID)
|
||||
|
||||
// ④ 等在跑 worker/planner drain(在跑 run 因夹逼的 MaxDuration 自行进收尾);
|
||||
// 超过 grace 仍未清空 → 硬 cancel 该任务 exec ctx(settling-aware 分支正确归类)。
|
||||
hardStop := time.Now().Add(settleDrainGrace)
|
||||
for e.inflightCount(t.ID) > 0 {
|
||||
if time.Now().After(hardStop) {
|
||||
log.Printf("[deadline] task %s drain 超时(%s),硬取消在跑 run", t.ID, settleDrainGrace)
|
||||
e.cancelExec(t.ID)
|
||||
_ = sleepCtx(ctx, 3*time.Second) // 给 worker 分支一点时间落库/归类
|
||||
break
|
||||
}
|
||||
if sleepCtx(ctx, 500*time.Millisecond) {
|
||||
return // 引擎整体关停
|
||||
}
|
||||
}
|
||||
|
||||
// ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。
|
||||
met := e.runFinalPlannerRound(ctx, t)
|
||||
|
||||
// ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done,
|
||||
// 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。
|
||||
status := "timeout"
|
||||
if met {
|
||||
status = "done"
|
||||
}
|
||||
won, err := e.m.SetTaskStatusGuarded(t.ID, status)
|
||||
switch {
|
||||
case err != nil:
|
||||
log.Printf("[deadline] task %s 落终态失败: %v", t.ID, err)
|
||||
case won:
|
||||
log.Printf("[deadline] task %s 收尾完成,终态=%s", t.ID, status)
|
||||
default:
|
||||
log.Printf("[deadline] task %s 收尾时已是终态,保留原状态", t.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// runFinalPlannerRound drives exactly ONE terminal planner round with the
|
||||
// task-timeout planner words (final goal judgment; no new intents). Waits for the
|
||||
// LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout.
|
||||
func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) {
|
||||
planner, _ := e.snapshotFor(t)
|
||||
for planner == nil {
|
||||
if sleepCtx(ctx, deadlinePollInterval) {
|
||||
return false
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
return false
|
||||
}
|
||||
planner, _ = e.snapshotFor(t)
|
||||
}
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))})
|
||||
// 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。
|
||||
fctx := e.clockCtx(ctx, t, true)
|
||||
e.incInflight(t.ID)
|
||||
tTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
met, reason, err := planner.Plan(fctx, tTaskID, e.m.assets, t.Store, t.Goal, t.drainDone(), emit)
|
||||
e.decInflight(t.ID)
|
||||
if err != nil {
|
||||
log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err)
|
||||
} else if met {
|
||||
log.Printf("[deadline] task %s 终局判定目标达成: %s", t.ID, reason)
|
||||
}
|
||||
return met
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
type goalSpec struct {
|
||||
Text string
|
||||
VulnClass string
|
||||
}
|
||||
|
||||
// createGoals materializes the goal node(s) under the task root (rel objective).
|
||||
// Decomposition is done ENTIRELY by the LLM (the project requires an LLM). There is
|
||||
// no rule-based fallback splitter — it only ever produced garbage (shredded URLs,
|
||||
// meaningless 2-way splits). If the LLM yields nothing (an error), the raw task goal
|
||||
// is used verbatim as a single goal so the task still has something to judge against.
|
||||
// Returns the seeded specs so callers can emit activity records for them.
|
||||
// emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI.
|
||||
func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec {
|
||||
// Use the SAME LLM the task runs on (its pinned profile, else the active profile),
|
||||
// NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars,
|
||||
// so FromEnv returned empty and every task silently fell back to the crude rule
|
||||
// splitter (which shredded URLs / made meaningless 2-way splits).
|
||||
var specs []goalSpec
|
||||
if cfg, ok := s.goalsLLMConfig(t); ok {
|
||||
for _, g := range agent.DecomposeGoals(ctx, cfg, t.Goal, t.Description, emit) {
|
||||
if strings.TrimSpace(g.Text) != "" {
|
||||
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(specs) == 0 {
|
||||
// No rule-based splitting: use the raw task goal verbatim as a single goal
|
||||
// (only reached on an LLM error — normal runs get real decomposed goals).
|
||||
if g := strings.TrimSpace(t.Goal); g != "" {
|
||||
log.Printf("[goals] task %s: LLM 目标拆解无产出，回退为「原始目标作为单目标」", t.ID)
|
||||
specs = []goalSpec{{Text: g}}
|
||||
}
|
||||
}
|
||||
origin, _ := t.Store.OriginFactID()
|
||||
for _, g := range specs {
|
||||
payload := map[string]any{"text": g.Text}
|
||||
if g.VulnClass != "" {
|
||||
payload["vulnclass"] = g.VulnClass
|
||||
}
|
||||
id, _ := t.Store.AddNode(db.KindGoal, payload, 0, "open", "system", nil)
|
||||
if origin > 0 && id > 0 {
|
||||
_ = t.Store.Link(origin, db.RelSpawns, id) // goals descend from the task root (origin fact)
|
||||
}
|
||||
}
|
||||
return specs
|
||||
}
|
||||
|
||||
// goalsLLMConfig resolves the LLM config for goal decomposition: the task's pinned
|
||||
// profile if any, else the global active profile (same source the engine runs on).
|
||||
func (s *Server) goalsLLMConfig(t *Task) (agent.Config, bool) {
|
||||
if t != nil && t.LLMProfileID != nil {
|
||||
if cfg, ok := s.loadProfileConfig(*t.LLMProfileID); ok {
|
||||
return cfg, true
|
||||
}
|
||||
}
|
||||
return s.loadLLMConfig()
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
)
|
||||
|
||||
// chatGuard returns a guard wired with the manager's interceptor, used for chat
|
||||
// conversations. Called once per applyLLM so a new LLM config always gets a fresh guard.
|
||||
func (s *Server) chatGuard() *guard.Guard {
|
||||
return guard.NewWithInterceptor(s.m.interceptor)
|
||||
}
|
||||
|
||||
// --- intercept rule CRUD ---
|
||||
|
||||
func (s *Server) interceptListRules(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
rules, err := pg.ListInterceptRules()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if rules == nil {
|
||||
rules = []db.InterceptRule{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"rules": rules})
|
||||
}
|
||||
|
||||
func (s *Server) interceptCreateRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req interceptRuleReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateInterceptRuleReq(req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
rule, err := pg.CreateInterceptRule(req.Name, req.MatchTarget, req.MatchType, req.Pattern, req.Action, req.Message, req.Priority, req.Enabled, req.TimeoutEnabled, req.TimeoutSeconds, req.TimeoutAction)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, rule)
|
||||
}
|
||||
|
||||
func (s *Server) interceptUpdateRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
var req interceptRuleReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateInterceptRuleReq(req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
rule, err := pg.UpdateInterceptRule(id, req.Name, req.MatchTarget, req.MatchType, req.Pattern, req.Action, req.Message, req.Priority, req.Enabled, req.TimeoutEnabled, req.TimeoutSeconds, req.TimeoutAction)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, rule)
|
||||
}
|
||||
|
||||
func (s *Server) interceptDeleteRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteInterceptRule(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, map[string]any{"deleted": id})
|
||||
}
|
||||
|
||||
func (s *Server) interceptToggleRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := pg.ToggleInterceptRule(id, req.Enabled); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, map[string]any{"ok": true, "enabled": req.Enabled})
|
||||
}
|
||||
|
||||
// --- pending (ask) ---
|
||||
|
||||
func (s *Server) interceptListPending(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
pending, err := pg.ListPendingIntercepts()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if pending == nil {
|
||||
pending = []db.InterceptPending{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"pending": pending})
|
||||
}
|
||||
|
||||
func (s *Server) interceptGetOne(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad pending id")
|
||||
return
|
||||
}
|
||||
p, err := pg.GetInterceptPending(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if p == nil {
|
||||
writeErr(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, p)
|
||||
}
|
||||
|
||||
func (s *Server) interceptListTaskItems(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
taskID := r.PathValue("taskID")
|
||||
if taskID == "" {
|
||||
writeErr(w, 400, "bad task id")
|
||||
return
|
||||
}
|
||||
items, err := pg.ListTaskIntercepts(taskID)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if items == nil {
|
||||
items = []db.InterceptApprovalRow{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) interceptHistory(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
items, err := pg.ListAllIntercepts(200)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if items == nil {
|
||||
items = []db.InterceptApprovalRow{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad pending id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Decision string `json:"decision"` // "allowed" | "denied"
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.Decision != "allowed" && req.Decision != "denied" {
|
||||
writeErr(w, 400, "decision 必须是 allowed 或 denied")
|
||||
return
|
||||
}
|
||||
if err := s.m.interceptor.Decide(id, req.Decision == "allowed"); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// --- tool-config (全局工具拦截范围) ---
|
||||
|
||||
// interceptGetToolConfig returns the list of tool names that are currently
|
||||
// configured to enter the intercept rule system.
|
||||
func (s *Server) interceptGetToolConfig(w http.ResponseWriter, r *http.Request) {
|
||||
tools, err := s.m.interceptor.GetEnabledTools()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"enabled_tools": tools})
|
||||
}
|
||||
|
||||
// interceptSetToolConfig replaces the list of tool names that should enter
|
||||
// the intercept rule system.
|
||||
func (s *Server) interceptSetToolConfig(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
EnabledTools []string `json:"enabled_tools"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.EnabledTools == nil {
|
||||
req.EnabledTools = []string{}
|
||||
}
|
||||
if err := s.m.interceptor.SetEnabledTools(req.EnabledTools); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
type interceptRuleReq struct {
|
||||
Name string `json:"name"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Priority int `json:"priority"`
|
||||
MatchTarget string `json:"match_target"`
|
||||
MatchType string `json:"match_type"`
|
||||
Pattern string `json:"pattern"`
|
||||
Action string `json:"action"`
|
||||
Message string `json:"message"`
|
||||
TimeoutEnabled bool `json:"timeout_enabled"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
TimeoutAction string `json:"timeout_action"`
|
||||
}
|
||||
|
||||
func validateInterceptRuleReq(req interceptRuleReq) error {
|
||||
if req.Name == "" {
|
||||
return fmt.Errorf("name 不能为空")
|
||||
}
|
||||
switch req.MatchTarget {
|
||||
case "tool_name", "tool_input":
|
||||
default:
|
||||
return fmt.Errorf("match_target 必须是 tool_name 或 tool_input")
|
||||
}
|
||||
switch req.MatchType {
|
||||
case "string", "regex":
|
||||
default:
|
||||
return fmt.Errorf("match_type 必须是 string 或 regex")
|
||||
}
|
||||
if req.Pattern == "" {
|
||||
return fmt.Errorf("pattern 不能为空")
|
||||
}
|
||||
switch req.Action {
|
||||
case "allow", "deny", "ask":
|
||||
default:
|
||||
return fmt.Errorf("action 必须是 allow、deny 或 ask")
|
||||
}
|
||||
if req.MatchType == "regex" {
|
||||
if _, err := regexp.Compile(req.Pattern); err != nil {
|
||||
return fmt.Errorf("pattern 不是有效正则：%w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// LogLine is one captured backend log entry exposed by the /api/logs endpoints.
|
||||
type LogLine struct {
|
||||
Seq int64 `json:"seq"`
|
||||
DBID int64 `json:"db_id,omitempty"` // server_logs.id; 0 for pre-persistence entries
|
||||
TS string `json:"ts"`
|
||||
Level string `json:"level"` // info | warn | error
|
||||
Tag string `json:"tag"` // the leading [tag] (pg / planner / activity / …), if any
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// dbWriteReq carries a log line to the async DB writer.
|
||||
type dbWriteReq struct {
|
||||
seq int64
|
||||
ll LogLine
|
||||
}
|
||||
|
||||
// logSinkT tees the standard logger into an in-memory ring buffer and fans new
|
||||
// lines out to SSE subscribers, so the UI can show a live backend log stream.
|
||||
// It still passes everything through to stderr (the terminal keeps working).
|
||||
type logSinkT struct {
|
||||
mu sync.Mutex
|
||||
ring []LogLine
|
||||
cap int
|
||||
seq int64
|
||||
subs map[chan LogLine]struct{}
|
||||
out io.Writer // passthrough (stderr)
|
||||
|
||||
dbOnce sync.Once
|
||||
dbCh chan dbWriteReq // buffered async channel; nil until SetDB is called
|
||||
}
|
||||
|
||||
var logSink = &logSinkT{cap: 3000, subs: map[chan LogLine]struct{}{}, out: os.Stderr}
|
||||
|
||||
// StartLogCapture redirects the standard log package through the in-memory sink
|
||||
// (still writing to stderr). Call once at startup, as early as possible.
|
||||
func StartLogCapture() { log.SetOutput(logSink) }
|
||||
|
||||
// SetDB wires a postgres DB into the sink once (idempotent). It:
|
||||
// 1. Restores the last 100 log rows from DB into the ring so the /logs page
|
||||
// shows history immediately after restart.
|
||||
// 2. Starts an async goroutine that persists every subsequent log line to DB.
|
||||
func (s *logSinkT) SetDB(ctx context.Context, pg *db.DB) {
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
s.dbOnce.Do(func() {
|
||||
ch := make(chan dbWriteReq, 2000)
|
||||
s.mu.Lock()
|
||||
s.dbCh = ch
|
||||
s.mu.Unlock()
|
||||
|
||||
// Restore last 100 rows from DB → prepend to ring as history context.
|
||||
if logs, err := pg.RecentLogs(100); err == nil && len(logs) > 0 {
|
||||
s.mu.Lock()
|
||||
restored := make([]LogLine, 0, len(logs))
|
||||
for _, l := range logs {
|
||||
s.seq++
|
||||
restored = append(restored, LogLine{
|
||||
Seq: s.seq,
|
||||
DBID: l.ID,
|
||||
TS: l.CreatedAt.Format(time.RFC3339),
|
||||
Level: l.Level,
|
||||
Tag: l.Tag,
|
||||
Text: l.Text,
|
||||
})
|
||||
}
|
||||
// Prepend history before any in-memory startup logs already in ring.
|
||||
s.ring = append(restored, s.ring...)
|
||||
if len(s.ring) > s.cap {
|
||||
s.ring = s.ring[len(s.ring)-s.cap:]
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Async writer: picks from channel and inserts into server_logs.
|
||||
// Uses os.Stderr directly to report errors and avoids recursive log calls.
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case req, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
id, err := pg.InsertLog(req.ll.Level, req.ll.Tag, req.ll.Text)
|
||||
if err != nil {
|
||||
_, _ = os.Stderr.Write([]byte("[logsink] db write: " + err.Error() + "\n"))
|
||||
continue
|
||||
}
|
||||
// Stamp DBID back into the ring entry so history pagination works.
|
||||
s.mu.Lock()
|
||||
for i := range s.ring {
|
||||
if s.ring[i].Seq == req.seq {
|
||||
s.ring[i].DBID = id
|
||||
break
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// Write implements io.Writer for the log package: one call per log.Printf line.
|
||||
func (s *logSinkT) Write(p []byte) (int, error) {
|
||||
_, _ = s.out.Write(p) // keep the terminal output
|
||||
line := strings.TrimRight(string(p), "\n")
|
||||
if strings.TrimSpace(line) != "" {
|
||||
s.add(parseLog(line))
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (s *logSinkT) add(l LogLine) {
|
||||
s.mu.Lock()
|
||||
s.seq++
|
||||
l.Seq = s.seq
|
||||
s.ring = append(s.ring, l)
|
||||
if len(s.ring) > s.cap {
|
||||
s.ring = s.ring[len(s.ring)-s.cap:]
|
||||
}
|
||||
ch := s.dbCh
|
||||
subs := make([]chan LogLine, 0, len(s.subs))
|
||||
for sub := range s.subs {
|
||||
subs = append(subs, sub)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// Async DB write (non-blocking; drops if channel is full under extreme load).
|
||||
if ch != nil {
|
||||
select {
|
||||
case ch <- dbWriteReq{seq: l.Seq, ll: l}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
for _, sub := range subs {
|
||||
select {
|
||||
case sub <- l:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// recent returns ring lines with Seq > since (capped to limit), and the latest seq.
|
||||
func (s *logSinkT) recent(since int64, limit int) (lines []LogLine, cursor int64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cursor = s.seq
|
||||
if limit <= 0 || limit > s.cap {
|
||||
limit = s.cap
|
||||
}
|
||||
for _, l := range s.ring {
|
||||
if l.Seq > since {
|
||||
lines = append(lines, l)
|
||||
}
|
||||
}
|
||||
if len(lines) > limit {
|
||||
lines = lines[len(lines)-limit:]
|
||||
}
|
||||
return lines, cursor
|
||||
}
|
||||
|
||||
func (s *logSinkT) subscribe() (<-chan LogLine, func()) {
|
||||
ch := make(chan LogLine, 256)
|
||||
s.mu.Lock()
|
||||
s.subs[ch] = struct{}{}
|
||||
s.mu.Unlock()
|
||||
var once sync.Once
|
||||
return ch, func() {
|
||||
once.Do(func() {
|
||||
s.mu.Lock()
|
||||
delete(s.subs, ch)
|
||||
s.mu.Unlock()
|
||||
close(ch)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// parseLog pulls a level + [tag] out of a standard-logger line; the timestamp is
|
||||
// the capture time (RFC3339).
|
||||
func parseLog(line string) LogLine {
|
||||
msg := line
|
||||
// strip the "2006/01/02 15:04:05" LstdFlags prefix if present
|
||||
if len(msg) >= 20 && msg[4] == '/' && msg[7] == '/' && msg[10] == ' ' {
|
||||
msg = strings.TrimSpace(msg[19:])
|
||||
}
|
||||
tag := ""
|
||||
if strings.HasPrefix(msg, "[") {
|
||||
if i := strings.IndexByte(msg, ']'); i > 1 {
|
||||
tag = msg[1:i]
|
||||
}
|
||||
}
|
||||
return LogLine{TS: time.Now().Format(time.RFC3339), Level: levelOf(msg), Tag: tag, Text: msg}
|
||||
}
|
||||
|
||||
func levelOf(msg string) string {
|
||||
low := strings.ToLower(msg)
|
||||
for _, k := range []string{"fatal", "panic", "error", "err:", "失败", "丢弃", "拒绝", "✕", "不可达"} {
|
||||
if strings.Contains(low, k) {
|
||||
return "error"
|
||||
}
|
||||
}
|
||||
for _, k := range []string{"warn", "disabled", "禁用", "skip", "stopped", "⚠", "重试"} {
|
||||
if strings.Contains(low, k) {
|
||||
return "warn"
|
||||
}
|
||||
}
|
||||
return "info"
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/mcp"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/mcphttp"
|
||||
)
|
||||
|
||||
// mcpClient is the shared surface of a connected MCP server (stdio or remote http),
|
||||
// so tools/list and cleanup are handled uniformly regardless of transport.
|
||||
type mcpClient interface {
|
||||
Tools(context.Context) ([]actool.CoreTool, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
// connectMCP dials one MCP server per its transport. Callers must Close the client.
|
||||
func connectMCP(ctx context.Context, m *db.MCPServer) (mcpClient, error) {
|
||||
switch m.Transport {
|
||||
case "stdio":
|
||||
if m.Command == "" {
|
||||
return nil, fmt.Errorf("stdio 传输缺少命令")
|
||||
}
|
||||
return mcp.NewStdioClient(ctx, m.Name, m.Command, jsonStrMap(m.Env), jsonStrSlice(m.Args)...)
|
||||
case "http":
|
||||
if m.URL == "" {
|
||||
return nil, fmt.Errorf("http 传输缺少 URL")
|
||||
}
|
||||
// env map doubles as HTTP headers (e.g. Authorization).
|
||||
return mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env))
|
||||
default:
|
||||
return nil, fmt.Errorf("未知传输方式 %q", m.Transport)
|
||||
}
|
||||
}
|
||||
|
||||
// discoverAndCacheMCP connects to one MCP, lists its tools, and persists the tool
|
||||
// names to mcp_tools_cache so the UI shows them without a live connection.
|
||||
func (s *Server) discoverAndCacheMCP(ctx context.Context, m *db.MCPServer) error {
|
||||
cl, err := connectMCP(ctx, m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cl.Close()
|
||||
ts, err := cl.Tools(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tools := make([]db.MCPTool, 0, len(ts))
|
||||
for _, t := range ts {
|
||||
tools = append(tools, db.MCPTool{Name: t.Name(), Description: t.Description()})
|
||||
}
|
||||
if err := s.m.pg.SaveMCPTools(m.ID, tools); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("[mcp] %s 发现 %d 个工具并已缓存", m.Name, len(tools))
|
||||
return nil
|
||||
}
|
||||
|
||||
// discoverEmptyMCPsOnStartup fills the tool cache for any enabled MCP that has none
|
||||
// yet (notably the seeded browser MCP on first run). Runs sequentially in one
|
||||
// goroutine so we never spawn many stdio servers (npx) at once, and never blocks
|
||||
// startup. Best-effort: a failure leaves the cache empty to retry next start.
|
||||
func (s *Server) discoverEmptyMCPsOnStartup() {
|
||||
servers, err := s.m.pg.ListMCP()
|
||||
if err != nil {
|
||||
log.Printf("[mcp] 启动自动发现: 读取列表失败: %v", err)
|
||||
return
|
||||
}
|
||||
for _, m := range servers {
|
||||
if !m.Enabled || len(m.Tools) > 0 {
|
||||
continue
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(s.ctx, 90*time.Second)
|
||||
if err := s.discoverAndCacheMCP(ctx, m); err != nil {
|
||||
log.Printf("[mcp] 启动自动发现 %s 失败: %v", m.Name, err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestMgmtAPI exercises the PostgreSQL-backed management API through the real mux.
|
||||
func TestMgmtAPI(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("database unavailable (%v) — skipping management API test", err)
|
||||
}
|
||||
if m.pg == nil {
|
||||
t.Skip("postgres unavailable — skipping management API test")
|
||||
}
|
||||
td := t.TempDir()
|
||||
s := New(context.Background(), m, td, td)
|
||||
h := s.Handler()
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatalf("sign jwt: %v", err)
|
||||
}
|
||||
|
||||
do := func(method, path string, body any) (int, map[string]any) {
|
||||
var r *http.Request
|
||||
if body != nil {
|
||||
b, _ := json.Marshal(body)
|
||||
r = httptest.NewRequest(method, path, bytes.NewReader(b))
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
r.Header.Set("Authorization", "Bearer "+tok)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, r)
|
||||
var out map[string]any
|
||||
json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
// agents list — must include at least the seeded builtins
|
||||
code, out := do("GET", "/api/agents", nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("GET agents: %d", code)
|
||||
}
|
||||
if ags, _ := out["agents"].([]any); len(ags) < 5 {
|
||||
t.Fatalf("want >=5 agents (seeded builtins), got %d", len(ags))
|
||||
}
|
||||
|
||||
// invalid template var → 400 (catalog whitelist)
|
||||
code, out = do("PUT", "/api/agents/planner/prompt", map[string]string{"template": "hi {{.Nope}}"})
|
||||
if code != 400 {
|
||||
t.Fatalf("bad var should be 400, got %d (%v)", code, out)
|
||||
}
|
||||
|
||||
// valid template using seeded catalog vars → 200
|
||||
code, _ = do("PUT", "/api/agents/planner/prompt", map[string]string{"template": "你是规划者，目标：{{.Goal}}，摘要 {{.AssetSummary}}"})
|
||||
if code != 200 {
|
||||
t.Fatalf("valid prompt save: %d", code)
|
||||
}
|
||||
|
||||
// preview renders with catalog examples
|
||||
code, out = do("POST", "/api/agents/planner/prompt/preview", map[string]any{})
|
||||
if code != 200 {
|
||||
t.Fatalf("preview: %d", code)
|
||||
}
|
||||
if rendered, _ := out["rendered"].(string); rendered == "" || rendered == "你是规划者，目标：{{.Goal}}，摘要 {{.AssetSummary}}" {
|
||||
t.Fatalf("preview did not substitute: %q", out["rendered"])
|
||||
}
|
||||
|
||||
// mcp create → visible to planner → resource-side sees planner → delete clears it
|
||||
// Remove any leftover from prior run to keep the test idempotent.
|
||||
m.pg.Exec(`DELETE FROM mcp_servers WHERE name = 't-itest'`)
|
||||
code, out = do("POST", "/api/mcp", map[string]any{"name": "t-itest", "transport": "stdio", "command": "x", "enabled": true})
|
||||
if code != 200 {
|
||||
t.Fatalf("create mcp: %d", code)
|
||||
}
|
||||
mid := int64(out["id"].(float64))
|
||||
|
||||
ag, _ := m.pg.GetAgentByKey("planner")
|
||||
code, _ = do("PUT", "/api/agents/planner/visibility", map[string]any{"mcp": []int64{mid}, "skill": []int64{}})
|
||||
if code != 200 {
|
||||
t.Fatalf("set visibility: %d", code)
|
||||
}
|
||||
code, out = do("GET", "/api/visibility/mcp/"+itoaTest(mid), nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("resource visibility: %d", code)
|
||||
}
|
||||
if agents, _ := out["agents"].([]any); len(agents) != 1 || agents[0].(string) != itoaTest(ag.ID) {
|
||||
t.Fatalf("resource-side should list planner, got %v", out["agents"])
|
||||
}
|
||||
|
||||
// cleanup
|
||||
do("DELETE", "/api/mcp/"+itoaTest(mid), nil)
|
||||
m.pg.Exec(`DELETE FROM agent_prompts WHERE agent_id=$1`, ag.ID)
|
||||
m.pg.Exec(`UPDATE agents SET current_prompt_id=NULL WHERE id=$1`, ag.ID)
|
||||
}
|
||||
|
||||
func itoaTest(n int64) string {
|
||||
b := []byte{}
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
for n > 0 {
|
||||
b = append([]byte{byte('0' + n%10)}, b...)
|
||||
n /= 10
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
)
|
||||
|
||||
// jsonResult marshals v to a JSON tool result.
|
||||
func jsonResult(v any) (actool.Result, error) {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(string(b)), nil
|
||||
}
|
||||
|
||||
// 本文件实现 P2「跨任务编排工具集」(docs/跑分编排 §2 P2)。这些是 host 工具——需要
|
||||
// 访问 Manager(任意任务的 Store)、Engine(暂停)、以及建任务流程,所以住在 server 层。
|
||||
// 读类工具把「现有 per-task 工具」重定向到目标任务的 store 上跑(建一个临时 ToolSet
|
||||
// 并 Call 其对应工具),从而复用完全相同的逻辑;控制类(spawn/pause)直接调 Manager/Engine。
|
||||
// 它们像流量工具一样 seed 进 tools 表、按 agent 绑定(只绑给编排 agent 才可见)。
|
||||
|
||||
// hostTools is the runtime host-tool provider fed to ToolAugment: traffic tools
|
||||
// (gated by capture) + cross-task orchestration tools + user-defined custom tools.
|
||||
// The second return is the names of custom tools flagged `deferred` (schema
|
||||
// withheld, routed via SearchExtraTools/ExecuteExtraTool). Per-agent binding still
|
||||
// decides who actually sees any of them.
|
||||
//
|
||||
//nolint:unused // used as the hostTools provider in wireAgentAugment
|
||||
func (s *Server) hostTools() ([]actool.CoreTool, map[string][]string) {
|
||||
tools := append(s.m.HostTools(), s.orchestrationTools()...)
|
||||
tools = append(tools, s.platformTools()...) // 平台操作工具(建改 skill/工具/MCP，给 Auto 用)
|
||||
custom, err := s.customTools()
|
||||
if err != nil {
|
||||
log.Printf("[custom-tool] 加载失败: %v", err)
|
||||
return tools, nil
|
||||
}
|
||||
tools = append(tools, custom...)
|
||||
// deferred custom tools → name -> its bound agent keys. ToolAugment turns a
|
||||
// name into a deferred entry only for agents it's actually bound to (so we don't
|
||||
// advertise a tool the per-agent binding will drop from the callable set).
|
||||
deferred := map[string][]string{}
|
||||
rows, _ := s.m.pg.ListCustomTools()
|
||||
for _, t := range rows {
|
||||
if t.Deferred && t.Enabled {
|
||||
deferred[t.Key] = t.Agents
|
||||
}
|
||||
}
|
||||
return tools, deferred
|
||||
}
|
||||
|
||||
// orchestrationTools returns the cross-task tool set. Bound per-agent via the
|
||||
// tools table (default: no binding — opt-in for orchestration agents).
|
||||
func (s *Server) orchestrationTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
s.toolListTasks(),
|
||||
s.toolListLLMProfiles(),
|
||||
s.toolSpawnTask(),
|
||||
s.toolPauseTask(),
|
||||
s.toolGetTaskGraph(),
|
||||
s.toolListTaskFindings(),
|
||||
s.toolAddHint(),
|
||||
s.toolGetWorkerTrace(),
|
||||
s.toolListWorkerTraces(),
|
||||
s.toolSearchWorkerTraces(),
|
||||
}
|
||||
}
|
||||
|
||||
// --- schema helpers ---
|
||||
|
||||
func strParam(desc string) map[string]any {
|
||||
return map[string]any{"type": "string", "description": desc}
|
||||
}
|
||||
|
||||
// parseProfileID reads an LLM profile id from a tool arg that may arrive as a JSON
|
||||
// number (5) or a numeric string ("5"); returns 0 when absent/unparseable.
|
||||
func parseProfileID(raw json.RawMessage) int64 {
|
||||
if len(raw) == 0 {
|
||||
return 0
|
||||
}
|
||||
var n int64
|
||||
if json.Unmarshal(raw, &n) == nil {
|
||||
return n
|
||||
}
|
||||
var str string
|
||||
if json.Unmarshal(raw, &str) == nil {
|
||||
v, _ := strconv.ParseInt(strings.TrimSpace(str), 10, 64)
|
||||
return v
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func objSchema(props map[string]any, required ...string) map[string]any {
|
||||
m := map[string]any{"type": "object", "properties": props}
|
||||
if len(required) > 0 {
|
||||
req := make([]any, len(required))
|
||||
for i, r := range required {
|
||||
req[i] = r
|
||||
}
|
||||
m["required"] = req
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func roTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||||
return actool.Build(actool.Spec{
|
||||
Name: name, Description: desc, Schema: schema,
|
||||
ReadOnly: func(json.RawMessage) bool { return true },
|
||||
Concurrent: func(json.RawMessage) bool { return true },
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
return run(ctx, in)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func wrTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||||
return actool.Build(actool.Spec{
|
||||
Name: name, Description: desc, Schema: schema,
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
return run(ctx, in)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// delegateToTask resolves the `task_id` in the input, builds a ToolSet bound to
|
||||
// that task's store, strips task_id, and calls the chosen per-task tool — so the
|
||||
// cross-task read reuses the exact in-task logic against another task.
|
||||
func (s *Server) delegateToTask(ctx context.Context, in json.RawMessage, pick func(*agent.ToolSet) actool.CoreTool) (actool.Result, error) {
|
||||
var head struct {
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &head)
|
||||
if strings.TrimSpace(head.TaskID) == "" {
|
||||
return actool.Errorf("task_id 为必填"), nil
|
||||
}
|
||||
t, ok := s.m.Task(head.TaskID)
|
||||
if !ok {
|
||||
return actool.Errorf("task 不存在: " + head.TaskID), nil
|
||||
}
|
||||
var m map[string]json.RawMessage
|
||||
_ = json.Unmarshal(in, &m)
|
||||
delete(m, "task_id")
|
||||
inner, _ := json.Marshal(m)
|
||||
tsx := agent.NewToolSet(t.Store, "orchestrator")
|
||||
if s.m.Assets() != nil {
|
||||
tsx.SetAssetStore(s.m.Assets(), s.m.Assets().Companies())
|
||||
}
|
||||
tsx.SetNotify(t.Notify) // hint writes wake this task's planner (no-op for read tools)
|
||||
return pick(tsx).Call(ctx, inner, nil)
|
||||
}
|
||||
|
||||
// --- tools ---
|
||||
|
||||
func (s *Server) toolListTasks() actool.CoreTool {
|
||||
return roTool("list_tasks",
|
||||
"列出所有任务(id/描述/目标/状态/运行时长/父任务/LLM 配置)，编排 agent 用它掌握全局、看哪些任务卡太久、各自用哪个 LLM。运行时长：运行中=创建→现在，终态=创建→最后活动(秒)。llm_profile：任务 planner/worker 用的配置名，(激活配置)=跟随全局激活。",
|
||||
objSchema(map[string]any{}),
|
||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||
lastAct, _ := s.m.PG().LastActivityAll()
|
||||
// id -> name to resolve each task's pinned LLM profile.
|
||||
profName := map[int64]string{}
|
||||
if profs, err := s.m.pg.ListProfiles(); err == nil {
|
||||
for _, p := range profs {
|
||||
profName[p.ID] = p.Name
|
||||
}
|
||||
}
|
||||
out := make([]map[string]any, 0)
|
||||
for _, t := range s.m.List() {
|
||||
status := s.deriveTaskStatus(t)
|
||||
end := lastAct[t.ExpID]
|
||||
if live := s.engine.LastActivity(t.ID); live > end {
|
||||
end = live
|
||||
}
|
||||
dur := int64(0)
|
||||
if status == "running" {
|
||||
dur = time.Now().Unix() - t.CreatedAt
|
||||
} else if end > t.CreatedAt {
|
||||
dur = end - t.CreatedAt
|
||||
}
|
||||
row := map[string]any{"id": t.ID, "description": t.Description, "goal": t.Goal, "status": status, "run_seconds": dur}
|
||||
if t.ParentRef != "" {
|
||||
row["parent_ref"] = t.ParentRef
|
||||
}
|
||||
if t.LLMProfileID == nil {
|
||||
row["llm_profile"] = "(激活配置)"
|
||||
} else if n, ok := profName[*t.LLMProfileID]; ok {
|
||||
row["llm_profile"] = n
|
||||
} else {
|
||||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return jsonResult(out)
|
||||
})
|
||||
}
|
||||
|
||||
// toolListLLMProfiles lists the available LLM profiles (name/model/active) so an
|
||||
// orchestration agent can pick one for spawn_task's llm_profile. Never leaks keys.
|
||||
func (s *Server) toolListLLMProfiles() actool.CoreTool {
|
||||
return roTool("list_llm_profiles",
|
||||
"列出可用的 LLM 配置(profile)：id、名称、模型、格式、是否为当前激活配置。用 id 给 spawn_task 的 llm_profile_id 参数指定子任务专属 LLM（如侦察用便宜模型、利用用强模型）。不含 API Key。",
|
||||
objSchema(map[string]any{}),
|
||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||
profs, err := s.m.pg.ListProfiles()
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
out := make([]map[string]any, 0, len(profs))
|
||||
for _, p := range profs {
|
||||
out = append(out, map[string]any{
|
||||
"id": p.ID, "name": p.Name, "model": p.Model, "format": p.Format, "is_active": p.IsDefault,
|
||||
})
|
||||
}
|
||||
return jsonResult(map[string]any{"profiles": out})
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolSpawnTask() actool.CoreTool {
|
||||
return wrTool("spawn_task",
|
||||
"新建一个子任务并启动探索引擎，返回 task_id。用于把一件事(如一道题/一个目标)派成独立任务。parent_ref 可选：填当前编排关联的父任务 id 做父子关联。",
|
||||
objSchema(map[string]any{
|
||||
"description": strParam("任务描述(简短标题)"),
|
||||
"goal": strParam("任务目标(要达成什么)"),
|
||||
"parent_ref": strParam("可选：父任务 id(做父子关联)"),
|
||||
"llm_profile_id": map[string]any{"type": "integer", "description": "可选：指定本子任务 planner/worker 用的 LLM 配置 id(见 list_llm_profiles)；留空则继承父任务、再回退全局激活配置"},
|
||||
"timeout_seconds": map[string]any{"type": "integer", "description": "可选：任务级超时(秒)。到点后触发优雅收尾并进入 timeout 终态；留空或 0 = 不限时"},
|
||||
}, "description", "goal"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
Description, Goal, ParentRef string
|
||||
LLMProfileID json.RawMessage `json:"llm_profile_id"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if strings.TrimSpace(a.Description) == "" {
|
||||
a.Description = "未命名任务"
|
||||
}
|
||||
if strings.TrimSpace(a.Goal) == "" {
|
||||
return actool.Errorf("goal 为必填"), nil
|
||||
}
|
||||
if a.TimeoutSeconds < 0 {
|
||||
a.TimeoutSeconds = 0
|
||||
}
|
||||
// LLM profile resolution: explicit id > inherit parent's pin > active(nil).
|
||||
var pin *int64
|
||||
if id := parseProfileID(a.LLMProfileID); id > 0 {
|
||||
if _, ok := s.loadProfileConfig(id); !ok {
|
||||
return actool.Errorf(fmt.Sprintf("LLM 配置 #%d 不存在或未设置 API Key", id)), nil
|
||||
}
|
||||
pin = &id
|
||||
} else if a.ParentRef != "" {
|
||||
if pt, ok := s.m.Task(a.ParentRef); ok {
|
||||
pin = pt.LLMProfileID
|
||||
}
|
||||
}
|
||||
t, err := s.m.CreateTask(a.Description, a.Goal, pin, a.TimeoutSeconds)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if a.ParentRef != "" {
|
||||
t.ParentRef = a.ParentRef
|
||||
if id, e := strconv.ParseInt(t.ID, 10, 64); e == nil {
|
||||
_ = s.m.PG().SetParentRef(id, a.ParentRef)
|
||||
}
|
||||
}
|
||||
s.seed(t, a.Description+" "+a.Goal) // seed 初始资产，喂给事件驱动 loop
|
||||
s.createGoals(s.ctx, t, nil) // 目标分解(LLM;规则兜底)
|
||||
s.engine.Run(s.ctx, t) // 启动该任务的探索引擎
|
||||
return actool.Text(fmt.Sprintf("task created: %s", t.ID)), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolPauseTask() actool.CoreTool {
|
||||
return wrTool("pause_task", "暂停指定任务(停止其 planner/worker 循环)。",
|
||||
objSchema(map[string]any{"task_id": strParam("要暂停的任务 id")}, "task_id"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if _, ok := s.m.Task(a.TaskID); !ok {
|
||||
return actool.Errorf("task 不存在: " + a.TaskID), nil
|
||||
}
|
||||
s.engine.Pause(a.TaskID)
|
||||
return actool.Text("task paused: " + a.TaskID), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolGetTaskGraph() actool.CoreTool {
|
||||
return roTool("get_task_graph", "读指定任务的探索图总览(同 graph_overview：资产计数/frontier/发现/覆盖等)，用 task_id 指定任务。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GraphOverviewTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolListTaskFindings() actool.CoreTool {
|
||||
return roTool("list_task_findings", "读指定任务的确认漏洞(含 flag/PoC；每条带 id/task_id/intent_id/vulnclass/severity/摘要/状态)，用 task_id 指定任务。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListFindingsTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolAddHint() actool.CoreTool {
|
||||
return wrTool("add_task_hint", "给指定任务注入战略提示(该任务的 planner 下轮生成意图时会读到)。\n"+
|
||||
"★优先批量：多条提示放进 hints 数组一次提交（返回 ids 数组，与 hints 等长同序，失败项 id=0）；单条则省略 hints 直接给顶层 text。",
|
||||
objSchema(map[string]any{
|
||||
"task_id": strParam("任务 id"),
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组，每个元素字段同顶层（text/asset_ids）。", "items": map[string]any{"type": "object"}},
|
||||
"text": strParam("[单条] 提示内容"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id（可选，0/1/多个；该任务内的资产 id）"},
|
||||
}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).AddHintTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolGetWorkerTrace() actool.CoreTool {
|
||||
return roTool("get_task_worker_trace",
|
||||
"看指定任务里某个 work(意图)的执行过程：get_task_worker_trace(task_id, intent_id) 看步骤摘要；再带 step_ids=[...] 取那几步完整内容(一次≤5)。",
|
||||
objSchema(map[string]any{
|
||||
"task_id": strParam("任务 id"),
|
||||
"intent_id": map[string]any{"type": "integer", "description": "意图 id(该任务里的 work)"},
|
||||
"step_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "可选：要取完整内容的步骤 id(≤5)"},
|
||||
}, "task_id", "intent_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GetWorkerTraceTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolListWorkerTraces() actool.CoreTool {
|
||||
return roTool("list_task_worker_traces", "列出指定任务里跑过哪些 work(意图) + 各自步数，用于发现哪些 work 值得翻看(再用 get_task_worker_trace)。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListWorkerTracesTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolSearchWorkerTraces() actool.CoreTool {
|
||||
return roTool("search_task_worker_traces", "在指定任务里按关键字搜索所有 work 的执行过程(返回命中步骤摘要 + intent_id)。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id"), "q": strParam("搜索关键字")}, "task_id", "q"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).SearchWorkerTracesTool)
|
||||
})
|
||||
}
|
||||
|
||||
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
|
||||
func (s *Server) deriveTaskStatus(t *Task) string {
|
||||
switch {
|
||||
case isTerminalStatus(t.Status):
|
||||
return t.Status
|
||||
case t.Paused || s.engine.IsPaused(t.ID):
|
||||
return "paused"
|
||||
case s.engine.Ready() && s.engine.Started(t.ID):
|
||||
return "running"
|
||||
}
|
||||
return "created"
|
||||
}
|
||||
|
||||
// orchestrationToolSeeds seeds the cross-task tools into the tools table so they
|
||||
// are bindable per-agent (default: bound to nobody — opt-in for orchestration
|
||||
// agents). First-insert only, like the traffic seeds.
|
||||
func (s *Server) seedOrchestrationTools() {
|
||||
// task-op + platform tools default-bind to the built-in Auto agent (它天生用来
|
||||
// 操作平台)。SeedTool 首插入生效;老库已 seed 的行由 seedAutoDefaultBindings 补绑。
|
||||
autoAgents, _ := json.Marshal([]string{"auto"})
|
||||
for _, t := range s.orchestrationTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||||
}
|
||||
for _, t := range s.platformTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||||
}
|
||||
s.refreshBuiltinToolSchemas()
|
||||
s.seedAutoDefaultBindings()
|
||||
s.seedPlannerDefaultBindings()
|
||||
s.seedAutoReportFindingBinding()
|
||||
// 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||||
// pentest 一起 seed 好了（项目尚无旧库，不做迁移）。
|
||||
}
|
||||
|
||||
// refreshBuiltinToolSchemas propagates code schema/description changes on the
|
||||
// orchestration + platform tools into already-seeded rows ONCE per version flag —
|
||||
// SeedTool is first-insert-only, so a new param (e.g. spawn_task 的 llm_profile) never
|
||||
// reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag.
|
||||
// Bump the flag whenever these tools' schemas/descriptions change in code.
|
||||
func (s *Server) refreshBuiltinToolSchemas() {
|
||||
const flag = "tool_schema_refresh_v4_goal_met_desc"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
tools := append(s.orchestrationTools(), s.platformTools()...)
|
||||
for _, t := range tools {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
if err := s.m.pg.RefreshToolDefaults(t.Name(), t.Description(), schema); err != nil {
|
||||
log.Printf("[tools] refresh %s schema failed: %v", t.Name(), err)
|
||||
}
|
||||
}
|
||||
// 同时把 planner 的 goal_met 描述刷成代码默认：旧库 seed 的描述带“结束本轮规划”的
|
||||
// 误导，会让 planner 把 goal_met 当成“结束空轮”的手段、刚开跑就误判整个任务完成。
|
||||
for _, sd := range agent.BuiltinToolSeeds() {
|
||||
if sd.Key != "goal_met" {
|
||||
continue
|
||||
}
|
||||
schema, _ := json.Marshal(sd.Schema)
|
||||
if err := s.m.pg.RefreshToolDefaults(sd.Key, sd.Desc, schema); err != nil {
|
||||
log.Printf("[tools] refresh %s desc failed: %v", sd.Key, err)
|
||||
}
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
log.Printf("[tools] 已刷新 orchestration/platform 工具 schema 到代码默认(一次性)")
|
||||
}
|
||||
|
||||
// seedAutoReportFindingBinding adds "auto" to report_finding's binding ONCE so
|
||||
// conversation-context agents can call it without requiring an intent_id.
|
||||
func (s *Server) seedAutoReportFindingBinding() {
|
||||
const flag = "auto_report_finding_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("auto", []string{"report_finding"}); err != nil {
|
||||
log.Printf("[auto] report_finding 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
// seedPlannerDefaultBindings adds "planner" to report_finding's binding ONCE
|
||||
// (guarded by a settings flag), so existing DBs — whose report_finding row was
|
||||
// seeded as worker-only — also let the planner record findings. Fresh DBs already
|
||||
// get it via PlannerTools(); this only backfills without overriding a user unbind.
|
||||
func (s *Server) seedPlannerDefaultBindings() {
|
||||
const flag = "planner_report_finding_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("planner", []string{"report_finding"}); err != nil {
|
||||
log.Printf("[planner] report_finding 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
// seedAutoDefaultBindings adds "auto" to the task-op + platform tools' bindings
|
||||
// ONCE (guarded by a settings flag), so existing DBs whose tool rows were seeded
|
||||
// before Auto existed still give Auto its default toolset — without re-adding it
|
||||
// after a user deliberately unbinds.
|
||||
func (s *Server) seedAutoDefaultBindings() {
|
||||
const flag = "auto_default_bindings_v3" // v3: 替换旧资产工具名，加入 insert_assets/add_company_scope
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
keys := make([]string, 0, len(platformToolKeys)+12)
|
||||
for _, t := range s.orchestrationTools() {
|
||||
keys = append(keys, t.Name())
|
||||
}
|
||||
keys = append(keys, platformToolKeys...)
|
||||
// 资产工具：Auto 操作平台常要看/登记资产、管理公司范围。
|
||||
keys = append(keys, "insert_assets", "add_company_scope", "list_assets")
|
||||
if err := s.m.pg.AddAgentToToolBinding("auto", keys); err != nil {
|
||||
log.Printf("[auto] 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 平台操作工具(给内置 Auto agent 用):建/改 skill、自定义工具、MCP。都是 host 工具,
|
||||
// seed 进 tools 表、默认绑定 auto,经 hostTools 注入。复用现有 db/文件系统逻辑。
|
||||
|
||||
func (s *Server) platformTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
s.toolCreateSkill(),
|
||||
s.toolUpdateSkillFile(),
|
||||
s.toolCreateCustomTool(),
|
||||
s.toolUpdateCustomTool(),
|
||||
s.toolCreateMCP(),
|
||||
s.toolUpdateMCP(),
|
||||
}
|
||||
}
|
||||
|
||||
// platformToolKeys are the tool keys the Auto agent gets bound by default.
|
||||
var platformToolKeys = []string{
|
||||
"create_skill", "update_skill_file",
|
||||
"create_custom_tool", "update_custom_tool",
|
||||
"create_mcp", "update_mcp",
|
||||
}
|
||||
|
||||
// ---- skills ----
|
||||
|
||||
func (s *Server) toolCreateSkill() actool.CoreTool {
|
||||
return wrTool("create_skill",
|
||||
"创建一个新 skill(写 SKILL.md，agentskills.io 规范)。name 小写字母/数字/连字符。",
|
||||
objSchema(map[string]any{
|
||||
"name": strParam("skill 名(小写字母开头，字母/数字/连字符)"),
|
||||
"description": strParam("skill 描述(必填，说明它做什么/何时用)"),
|
||||
"instructions": strParam("Markdown 正文说明(可选)"),
|
||||
}, "name", "description"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct{ Name, Description, Instructions string }
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if !validSkillName(a.Name) {
|
||||
return actool.Errorf("skill 名不合法(小写字母开头，仅字母/数字/连字符，≤64)"), nil
|
||||
}
|
||||
if strings.TrimSpace(a.Description) == "" {
|
||||
return actool.Errorf("description 必填"), nil
|
||||
}
|
||||
path := filepath.Join(s.skillDir, a.Name)
|
||||
if _, err := os.Stat(path); err == nil {
|
||||
return actool.Errorf("skill 已存在: " + a.Name), nil
|
||||
}
|
||||
if err := os.MkdirAll(path, 0o755); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("---\n")
|
||||
fmt.Fprintf(&b, "name: %s\n", a.Name)
|
||||
fmt.Fprintf(&b, "description: %s\n", a.Description)
|
||||
b.WriteString("---\n")
|
||||
if strings.TrimSpace(a.Instructions) != "" {
|
||||
b.WriteString(a.Instructions)
|
||||
} else {
|
||||
fmt.Fprintf(&b, "## %s\n\n1. \n2. \n3. \n", a.Name)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte(b.String()), 0o644); err != nil {
|
||||
_ = os.RemoveAll(path)
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("skill created: " + a.Name), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateSkillFile() actool.CoreTool {
|
||||
return wrTool("update_skill_file",
|
||||
"写/覆盖某个 skill 内的一个文件(默认 SKILL.md)。用于修改技能内容或加脚本/引用。",
|
||||
objSchema(map[string]any{
|
||||
"name": strParam("skill 名"),
|
||||
"file": strParam("相对路径(可选，默认 SKILL.md，如 scripts/run.py)"),
|
||||
"content": strParam("文件完整内容"),
|
||||
}, "name", "content"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct{ Name, File, Content string }
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if !validSkillName(a.Name) {
|
||||
return actool.Errorf("skill 名不合法"), nil
|
||||
}
|
||||
skillPath := filepath.Join(s.skillDir, a.Name)
|
||||
if _, err := os.Stat(skillPath); os.IsNotExist(err) {
|
||||
return actool.Errorf("skill 不存在: " + a.Name), nil
|
||||
}
|
||||
rel := strings.TrimSpace(a.File)
|
||||
if rel == "" {
|
||||
rel = "SKILL.md"
|
||||
}
|
||||
clean, msg := skillRelPath(rel)
|
||||
if msg != "" {
|
||||
return actool.Errorf("非法路径: " + msg), nil
|
||||
}
|
||||
full := filepath.Join(skillPath, clean)
|
||||
if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if err := os.WriteFile(full, []byte(a.Content), 0o644); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("skill file written: " + a.Name + "/" + clean), nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- custom tools ----
|
||||
|
||||
type customToolToolInput struct {
|
||||
Key string `json:"key"`
|
||||
Description string `json:"description"`
|
||||
Kind string `json:"kind"`
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Schema json.RawMessage `json:"schema"`
|
||||
Agents []string `json:"agents"`
|
||||
Deferred bool `json:"deferred"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func customToolSchema(keyDesc string) map[string]any {
|
||||
return objSchema(map[string]any{
|
||||
"key": strParam(keyDesc),
|
||||
"description": strParam("发给模型的描述"),
|
||||
"kind": strParam("shell | command | script(仅Python) | http。shell=bash 环境声明(仅告知模型该工具可在 bash 中直接调用，无需 exec/schema)；其余三种需提供 exec"),
|
||||
"exec": map[string]any{"type": "object", "description": "执行规格(shell 类型不需要): command→{command}; script→{code}; http→{method,url,headers,body,proxy,use_recording_proxy}"},
|
||||
"schema": map[string]any{"type": "object", "description": "参数 JSON-Schema(shell/command/script 可留空; http 必填且需含 properties)"},
|
||||
"agents": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "绑定的 agent key(可选)"},
|
||||
"deferred": map[string]any{"type": "boolean", "description": "是否延迟(shell 类型无效；仅 command/script/http 的不常用工具才开)"},
|
||||
"enabled": map[string]any{"type": "boolean", "description": "是否启用(默认 true)"},
|
||||
}, "key", "kind")
|
||||
}
|
||||
|
||||
func toDBTool(a customToolToolInput) *db.Tool {
|
||||
enabled := true
|
||||
if a.Enabled != nil {
|
||||
enabled = *a.Enabled
|
||||
}
|
||||
return &db.Tool{
|
||||
Key: a.Key, Description: a.Description, Schema: a.Schema, Agents: a.Agents,
|
||||
Enabled: enabled, Kind: a.Kind, Exec: a.Exec, Deferred: a.Deferred,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) toolCreateCustomTool() actool.CoreTool {
|
||||
return wrTool("create_custom_tool", "创建一个自定义工具(shell/command/script/http)。shell=bash 环境声明，只需 key+description+agents，无需 exec/schema。",
|
||||
customToolSchema("工具 key(小写字母开头，字母/数字/下划线)"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a customToolToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
a.Key = strings.TrimSpace(a.Key)
|
||||
if !reToolKey.MatchString(a.Key) {
|
||||
return actool.Errorf("key 需小写字母开头，仅含小写字母/数字/下划线"), nil
|
||||
}
|
||||
if a.Kind != "shell" && a.Kind != "command" && a.Kind != "script" && a.Kind != "http" {
|
||||
return actool.Errorf("kind 需为 shell / command / script / http"), nil
|
||||
}
|
||||
if a.Kind == "http" && !hasSchemaProps(a.Schema) {
|
||||
return actool.Errorf("http 工具必须提供参数 JSON Schema(不能留空)"), nil
|
||||
}
|
||||
if exist, _ := s.m.pg.GetTool(a.Key); exist != nil {
|
||||
return actool.Errorf("该 key 已存在: " + a.Key), nil
|
||||
}
|
||||
if err := s.m.pg.CreateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("custom tool created: " + a.Key), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateCustomTool() actool.CoreTool {
|
||||
return wrTool("update_custom_tool", "修改一个已有的自定义工具(按 key)。",
|
||||
customToolSchema("要修改的自定义工具 key"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a customToolToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
existing, _ := s.m.pg.GetTool(a.Key)
|
||||
if existing == nil || existing.System {
|
||||
return actool.Errorf("只能修改自定义工具: " + a.Key), nil
|
||||
}
|
||||
if a.Kind != "shell" && a.Kind != "command" && a.Kind != "script" && a.Kind != "http" {
|
||||
return actool.Errorf("kind 需为 shell / command / script / http"), nil
|
||||
}
|
||||
if a.Kind == "http" && !hasSchemaProps(a.Schema) {
|
||||
return actool.Errorf("http 工具必须提供参数 JSON Schema(不能留空)"), nil
|
||||
}
|
||||
if err := s.m.pg.UpdateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("custom tool updated: " + a.Key), nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- MCP ----
|
||||
|
||||
type mcpToolInput struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
Command string `json:"command"`
|
||||
Args json.RawMessage `json:"args"`
|
||||
Env json.RawMessage `json:"env"`
|
||||
URL string `json:"url"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func mcpSchema(withID bool) map[string]any {
|
||||
props := map[string]any{
|
||||
"name": strParam("MCP 服务器名"),
|
||||
"transport": strParam("stdio | http / sse"),
|
||||
"command": strParam("stdio 的启动命令(如 npx)"),
|
||||
"args": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "命令参数数组"},
|
||||
"env": map[string]any{"type": "object", "description": "环境变量 {KEY:VALUE}"},
|
||||
"url": strParam("http/sse 的 URL"),
|
||||
"enabled": map[string]any{"type": "boolean", "description": "是否启用(默认 true)"},
|
||||
}
|
||||
required := []string{"name", "transport"}
|
||||
if withID {
|
||||
props["id"] = map[string]any{"type": "integer", "description": "要修改的 MCP 服务器 id"}
|
||||
required = []string{"id", "name", "transport"}
|
||||
}
|
||||
return objSchema(props, required...)
|
||||
}
|
||||
|
||||
func (a mcpToolInput) toDB() *db.MCPServer {
|
||||
enabled := true
|
||||
if a.Enabled != nil {
|
||||
enabled = *a.Enabled
|
||||
}
|
||||
return &db.MCPServer{
|
||||
ID: a.ID, Name: a.Name, Transport: a.Transport, Command: a.Command,
|
||||
Args: a.Args, Env: a.Env, URL: a.URL, Enabled: enabled,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) toolCreateMCP() actool.CoreTool {
|
||||
return wrTool("create_mcp", "创建一个 MCP 服务器(stdio/http)。创建后其工具需按 agent 可见性授权。",
|
||||
mcpSchema(false),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a mcpToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
a.ID = 0
|
||||
if strings.TrimSpace(a.Name) == "" || strings.TrimSpace(a.Transport) == "" {
|
||||
return actool.Errorf("name / transport 必填"), nil
|
||||
}
|
||||
id, err := s.m.pg.SaveMCP(a.toDB())
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(fmt.Sprintf("mcp created: id=%d name=%s", id, a.Name)), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateMCP() actool.CoreTool {
|
||||
return wrTool("update_mcp", "修改一个已有的 MCP 服务器(按 id)。",
|
||||
mcpSchema(true),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a mcpToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if a.ID == 0 {
|
||||
return actool.Errorf("id 必填"), nil
|
||||
}
|
||||
if _, err := s.m.pg.SaveMCP(a.toDB()); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(fmt.Sprintf("mcp updated: id=%d", a.ID)), nil
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestSkillTools functionally exercises create_skill + update_skill_file against a
|
||||
// temp skill dir (no DB needed — these tools only touch the filesystem).
|
||||
func TestSkillTools(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s := &Server{skillDir: dir}
|
||||
|
||||
// create_skill writes a spec-compliant SKILL.md.
|
||||
in, _ := json.Marshal(map[string]any{
|
||||
"name": "my-skill", "description": "does X", "instructions": "## do\n1. a",
|
||||
})
|
||||
res, err := s.toolCreateSkill().Call(context.Background(), in, nil)
|
||||
if err != nil || res.IsError {
|
||||
t.Fatalf("create_skill: err=%v result=%+v", err, res)
|
||||
}
|
||||
md, err := os.ReadFile(filepath.Join(dir, "my-skill", "SKILL.md"))
|
||||
if err != nil {
|
||||
t.Fatalf("SKILL.md not written: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "name: my-skill") || !strings.Contains(string(md), "description: does X") || !strings.Contains(string(md), "1. a") {
|
||||
t.Fatalf("SKILL.md content:\n%s", md)
|
||||
}
|
||||
|
||||
// creating the same skill again fails.
|
||||
if r, _ := s.toolCreateSkill().Call(context.Background(), in, nil); !r.IsError {
|
||||
t.Fatal("duplicate create should error")
|
||||
}
|
||||
|
||||
// invalid name rejected.
|
||||
bad, _ := json.Marshal(map[string]any{"name": "Bad Name", "description": "x"})
|
||||
if r, _ := s.toolCreateSkill().Call(context.Background(), bad, nil); !r.IsError {
|
||||
t.Fatal("invalid skill name should error")
|
||||
}
|
||||
|
||||
// update_skill_file writes a nested script; path traversal is rejected.
|
||||
up, _ := json.Marshal(map[string]any{"name": "my-skill", "file": "scripts/run.py", "content": "print(1)"})
|
||||
if r, err := s.toolUpdateSkillFile().Call(context.Background(), up, nil); err != nil || r.IsError {
|
||||
t.Fatalf("update_skill_file: err=%v result=%+v", err, r)
|
||||
}
|
||||
got, err := os.ReadFile(filepath.Join(dir, "my-skill", "scripts", "run.py"))
|
||||
if err != nil || string(got) != "print(1)" {
|
||||
t.Fatalf("nested file not written: %v %q", err, got)
|
||||
}
|
||||
evil, _ := json.Marshal(map[string]any{"name": "my-skill", "file": "../escape.txt", "content": "x"})
|
||||
if r, _ := s.toolUpdateSkillFile().Call(context.Background(), evil, nil); !r.IsError {
|
||||
t.Fatal("path traversal should be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Scheduler drives P3 triggers: on each tick it fires due interval triggers and
|
||||
// scans for new findings / newly-met goals (any task) to fire event triggers.
|
||||
// Each fire enqueues a NEW conversation on the agent's per-agent trigger queue
|
||||
// (StartTriggeredRun): fires for the same agent run one at a time in FIFO order,
|
||||
// distinct agents still run concurrently. State is persisted (per-trigger last_fire
|
||||
// + finding watermark + fired-goal set) so a restart resumes without double-firing.
|
||||
// Triggers only attach to CUSTOM agents.
|
||||
type Scheduler struct {
|
||||
s *Server
|
||||
pg *db.DB
|
||||
tick time.Duration
|
||||
}
|
||||
|
||||
const (
|
||||
schedKeyLastFinding = "last_finding_id" // watermark: max finding node id fired for
|
||||
schedKeyFiredGoals = "fired_goals" // JSON array of goal node ids already fired
|
||||
schedKeyLastTimeout = "last_timeout_id" // watermark: max task id fired for task timeout
|
||||
)
|
||||
|
||||
func newScheduler(s *Server) *Scheduler {
|
||||
return &Scheduler{s: s, pg: s.m.pg, tick: 5 * time.Second}
|
||||
}
|
||||
|
||||
// Run loops until ctx is done, ticking the scheduler. Started once from server New.
|
||||
func (sc *Scheduler) Run(ctx context.Context) {
|
||||
if sc.pg == nil {
|
||||
return
|
||||
}
|
||||
sc.init()
|
||||
t := time.NewTicker(sc.tick)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
sc.step()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// init seeds the watermarks on first run so pre-existing findings/goals don't all
|
||||
// fire at once — only events created AFTER the scheduler first starts count.
|
||||
func (sc *Scheduler) init() {
|
||||
if sc.mustState(schedKeyLastFinding) == "" {
|
||||
var maxID int64
|
||||
if evs, err := sc.pg.NewFindingsSince(0); err == nil {
|
||||
for _, e := range evs {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastFinding, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
if sc.mustState(schedKeyFiredGoals) == "" {
|
||||
set := map[int64]bool{}
|
||||
if evs, err := sc.pg.MetGoals(); err == nil {
|
||||
for _, e := range evs {
|
||||
set[e.NodeID] = true
|
||||
}
|
||||
}
|
||||
sc.saveFiredGoalSet(set)
|
||||
}
|
||||
if sc.mustState(schedKeyLastTimeout) == "" {
|
||||
var maxID int64
|
||||
if evs, err := sc.pg.TimedOutTasksSince(0); err == nil {
|
||||
for _, e := range evs {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastTimeout, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
}
|
||||
|
||||
func (sc *Scheduler) step() {
|
||||
triggers, err := sc.pg.ListEnabledTriggers()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(triggers) == 0 {
|
||||
return
|
||||
}
|
||||
sc.fireIntervals(triggers)
|
||||
sc.fireFindings(triggers)
|
||||
sc.fireGoals(triggers)
|
||||
sc.fireTaskTimeouts(triggers)
|
||||
}
|
||||
|
||||
// fireIntervals fires triggers whose interval has elapsed since last_fire.
|
||||
func (sc *Scheduler) fireIntervals(triggers []*db.AgentTrigger) {
|
||||
now := time.Now()
|
||||
for _, tr := range triggers {
|
||||
if tr.IntervalSec <= 0 {
|
||||
continue
|
||||
}
|
||||
due := tr.LastFire == nil || now.Sub(*tr.LastFire) >= time.Duration(tr.IntervalSec)*time.Second
|
||||
if !due {
|
||||
continue
|
||||
}
|
||||
_ = sc.pg.TouchTriggerFire(tr.ID)
|
||||
ctx := "\n\n【本次为定时触发】" + now.Format(" 2006-01-02 15:04:05 MST")
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("定时触发 · %s", now.Format("15:04")), tr.IntervalMessage+ctx, 0, false)
|
||||
}
|
||||
}
|
||||
|
||||
// fireFindings fires on_finding triggers for findings above the persisted
|
||||
// watermark (monotonic node id → no double-fire across restarts).
|
||||
func (sc *Scheduler) fireFindings(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnFinding {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64)
|
||||
events, err := sc.pg.NewFindingsSince(last)
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
maxID := last
|
||||
for _, e := range events {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务发现 finding 触发】\n任务: #%d %s（目标：%s）\n发现: [%s/%s] %s",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal, e.VulnClass, e.Severity, e.Summary)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("finding 触发 · task#%d", e.TaskID), tr.FindingMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastFinding, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
|
||||
// fireGoals fires on_goal_met triggers for met goals not yet in the fired set.
|
||||
func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnGoalMet {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
events, err := sc.pg.MetGoals()
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
fired := sc.firedGoalSet()
|
||||
changed := false
|
||||
for _, e := range events {
|
||||
if fired[e.NodeID] {
|
||||
continue
|
||||
}
|
||||
fired[e.NodeID] = true
|
||||
changed = true
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s（目标：%s）\n达成目标: %s",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("目标触发 · task#%d", e.TaskID), tr.GoalMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
sc.saveFiredGoalSet(fired)
|
||||
}
|
||||
}
|
||||
|
||||
// fireTaskTimeouts fires on_task_timeout triggers for tasks that newly reached
|
||||
// status='timeout' above the persisted watermark (task id → no double-fire).
|
||||
func (sc *Scheduler) fireTaskTimeouts(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnTaskTimeout {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64)
|
||||
events, err := sc.pg.TimedOutTasksSince(last)
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
maxID := last
|
||||
for _, e := range events {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s（目标：%s）",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("超时触发 · task#%d", e.TaskID), tr.TaskTimeoutMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastTimeout, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
|
||||
func (sc *Scheduler) mustState(key string) string {
|
||||
v, _ := sc.pg.GetSchedState(key)
|
||||
return v
|
||||
}
|
||||
|
||||
func (sc *Scheduler) firedGoalSet() map[int64]bool {
|
||||
out := map[int64]bool{}
|
||||
raw := sc.mustState(schedKeyFiredGoals)
|
||||
if raw == "" {
|
||||
return out
|
||||
}
|
||||
var ids []int64
|
||||
if json.Unmarshal([]byte(raw), &ids) == nil {
|
||||
for _, id := range ids {
|
||||
out[id] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (sc *Scheduler) saveFiredGoalSet(set map[int64]bool) {
|
||||
ids := make([]int64, 0, len(set))
|
||||
for id := range set {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
b, _ := json.Marshal(ids)
|
||||
if err := sc.pg.SetSchedState(schedKeyFiredGoals, string(b)); err != nil {
|
||||
log.Printf("[scheduler] save fired goals failed: %v", err)
|
||||
}
|
||||
}
|
||||
+1757
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,513 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/mcphttp"
|
||||
)
|
||||
|
||||
// 资产同步（ScopeSentry 数据源）。
|
||||
//
|
||||
// ScopeSentry 是一个 ASM 资产测绘平台，通过其 MCP 接口按【项目】/【任务】两个维度
|
||||
// 拉取子域名、Web 应用、服务等资产，映射为 ARTEX 的公司 + 资产模型。数据源本身是
|
||||
// 一个名为 "ScopeSentry" 的 http 传输 MCP 行（url + X-API-Key 头保存在 mcp_servers）。
|
||||
//
|
||||
// 与 agent 工具层不同，这里用 mcphttp.Client.Call 直接调用 MCP 工具、拿原始 JSON，
|
||||
// 不触发 AskUser 权限弹窗（后台批量同步）。
|
||||
|
||||
const (
|
||||
scopeSentryMCPName = "ScopeSentry"
|
||||
syncMaxPerType = 5000 // 单目标单类型的入库保护上限
|
||||
syncDefaultPage = 100
|
||||
)
|
||||
|
||||
// findMCPByName returns the MCP server row with the given name, or nil.
|
||||
func (s *Server) findMCPByName(name string) (*db.MCPServer, error) {
|
||||
all, err := s.m.pg.ListMCP()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, m := range all {
|
||||
if m.Name == name {
|
||||
return m, nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// scopeSentryClient dials the configured ScopeSentry MCP (http transport). The env
|
||||
// map doubles as HTTP headers (X-API-Key). Callers must Close the client.
|
||||
func (s *Server) scopeSentryClient(ctx context.Context) (*mcphttp.Client, error) {
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if m == nil {
|
||||
return nil, fmt.Errorf("数据源 %s 不存在，请先创建", scopeSentryMCPName)
|
||||
}
|
||||
if m.URL == "" {
|
||||
return nil, fmt.Errorf("数据源 %s 未配置 URL，请先配置", scopeSentryMCPName)
|
||||
}
|
||||
return mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env))
|
||||
}
|
||||
|
||||
// envHasValue reports whether the env/header map has any non-empty value
|
||||
// (i.e. an API key / Authorization header was filled in).
|
||||
func envHasValue(raw json.RawMessage) bool {
|
||||
for _, v := range jsonStrMap(raw) {
|
||||
if strings.TrimSpace(v) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/status ----------
|
||||
|
||||
func (s *Server) syncSSStatus(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
resp := map[string]any{
|
||||
"exists": m != nil,
|
||||
"configured": false,
|
||||
"enabled": false,
|
||||
"reachable": false,
|
||||
"tools": []string{},
|
||||
}
|
||||
if m == nil {
|
||||
writeJSON(w, 200, resp)
|
||||
return
|
||||
}
|
||||
configured := m.URL != "" && envHasValue(m.Env)
|
||||
resp["configured"] = configured
|
||||
resp["enabled"] = m.Enabled
|
||||
resp["url"] = m.URL
|
||||
if m.Tools != nil {
|
||||
resp["tools"] = m.Tools
|
||||
}
|
||||
// Light reachability probe only when it can actually connect.
|
||||
if configured && m.Enabled {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 12*time.Second)
|
||||
defer cancel()
|
||||
if cl, cerr := mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env)); cerr == nil {
|
||||
if _, terr := cl.Tools(ctx); terr == nil {
|
||||
resp["reachable"] = true
|
||||
}
|
||||
_ = cl.Close()
|
||||
}
|
||||
}
|
||||
writeJSON(w, 200, resp)
|
||||
}
|
||||
|
||||
// ---------- POST /api/sync/scopesentry/datasource ----------
|
||||
// Creates the placeholder row if missing, and/or sets URL + API key and enables it.
|
||||
|
||||
func (s *Server) syncSSDatasource(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
URL string `json:"url"`
|
||||
APIKey string `json:"api_key"`
|
||||
}
|
||||
_ = decode(r, &body) // empty body = just create the placeholder
|
||||
body.URL = strings.TrimSpace(body.URL)
|
||||
body.APIKey = strings.TrimSpace(body.APIKey)
|
||||
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if m == nil {
|
||||
m = &db.MCPServer{Name: scopeSentryMCPName, Transport: "http", Env: json.RawMessage(`{"X-API-Key":""}`)}
|
||||
}
|
||||
if body.URL != "" {
|
||||
m.URL = body.URL
|
||||
}
|
||||
if body.APIKey != "" {
|
||||
env, _ := json.Marshal(map[string]string{"X-API-Key": body.APIKey})
|
||||
m.Env = env
|
||||
}
|
||||
// Enable only once it can actually be used.
|
||||
m.Enabled = m.URL != "" && envHasValue(m.Env)
|
||||
|
||||
id, err := pg.SaveMCP(m)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
m.ID = id
|
||||
// Best-effort tool discovery so the status card shows tools right away.
|
||||
if m.Enabled {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
|
||||
_ = s.discoverAndCacheMCP(ctx, m)
|
||||
cancel()
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"id": id, "enabled": m.Enabled})
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/projects ----------
|
||||
|
||||
func (s *Server) syncSSProjects(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
args := map[string]any{
|
||||
"pageIndex": queryInt(r, "page", 1),
|
||||
"pageSize": queryInt(r, "size", 50),
|
||||
}
|
||||
if q := r.URL.Query().Get("search"); q != "" {
|
||||
args["search"] = q
|
||||
}
|
||||
text, err := cl.Call(ctx, "list_projects_data", args)
|
||||
if err != nil {
|
||||
writeErr(w, 502, "list_projects_data 失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
// {result:{All:[{id,name,logo,AssetCount,tag}], <tag>:[...]}, tag:{...}}
|
||||
var env struct {
|
||||
Result map[string]json.RawMessage `json:"result"`
|
||||
Tag map[string]int `json:"tag"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
writeErr(w, 502, "解析项目列表失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
projects := json.RawMessage("[]")
|
||||
if raw, ok := env.Result["All"]; ok && len(raw) > 0 {
|
||||
projects = raw
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"projects": projects, "tag": env.Tag})
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/tasks ----------
|
||||
|
||||
func (s *Server) syncSSTasks(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
args := map[string]any{
|
||||
"pageIndex": queryInt(r, "page", 1),
|
||||
"pageSize": queryInt(r, "size", 50),
|
||||
}
|
||||
if q := r.URL.Query().Get("search"); q != "" {
|
||||
args["search"] = q
|
||||
}
|
||||
text, err := cl.Call(ctx, "list_tasks", args)
|
||||
if err != nil {
|
||||
writeErr(w, 502, "list_tasks 失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
var env struct {
|
||||
List json.RawMessage `json:"list"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
writeErr(w, 502, "解析任务列表失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
tasks := env.List
|
||||
if len(tasks) == 0 {
|
||||
tasks = json.RawMessage("[]")
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"tasks": tasks})
|
||||
}
|
||||
|
||||
// ---------- POST /api/sync/scopesentry/sync ----------
|
||||
|
||||
type ssSyncReq struct {
|
||||
Dimension string `json:"dimension"` // "project" | "task"
|
||||
Targets []string `json:"targets"` // project ObjectIDs, or task names
|
||||
AssetTypes []string `json:"asset_types"` // subdomain | app | service
|
||||
CreateCompany bool `json:"create_company"`
|
||||
PageSize int `json:"page_size"`
|
||||
}
|
||||
|
||||
func (s *Server) syncSSRun(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
as := s.assetStore()
|
||||
cs := s.companyStore()
|
||||
if as == nil || cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req ssSyncReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.Dimension != "project" && req.Dimension != "task" {
|
||||
writeErr(w, 400, "dimension 必须是 project 或 task")
|
||||
return
|
||||
}
|
||||
if len(req.Targets) == 0 {
|
||||
writeErr(w, 400, "targets 不能为空")
|
||||
return
|
||||
}
|
||||
if len(req.AssetTypes) == 0 {
|
||||
req.AssetTypes = []string{"subdomain", "service", "app"}
|
||||
}
|
||||
pageSize := req.PageSize
|
||||
if pageSize <= 0 || pageSize > 500 {
|
||||
pageSize = syncDefaultPage
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
synced := map[string]int{"subdomain": 0, "app": 0, "service": 0, "ip": 0}
|
||||
var companies []string
|
||||
var warnings, errs []string
|
||||
madeCompany := false
|
||||
|
||||
for _, target := range req.Targets {
|
||||
filter := map[string]any{}
|
||||
switch req.Dimension {
|
||||
case "project":
|
||||
filter["project"] = []string{target}
|
||||
if req.CreateCompany {
|
||||
name, roots, perr := s.ssProjectMeta(ctx, cl, target)
|
||||
if perr != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("项目 %s 详情获取失败: %v", target, perr))
|
||||
} else if name != "" {
|
||||
cid, cerr := cs.UpsertByName(name)
|
||||
if cerr != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("公司 %s 创建失败: %v", name, cerr))
|
||||
} else {
|
||||
companies = append(companies, name)
|
||||
madeCompany = true
|
||||
if len(roots) > 0 {
|
||||
cs.AddScope(cid, roots, "scopesentry:"+name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
case "task":
|
||||
filter["task"] = []string{target}
|
||||
}
|
||||
|
||||
for _, at := range req.AssetTypes {
|
||||
ssType, ok := map[string]string{"subdomain": "subdomain", "service": "asset", "app": "app"}[at]
|
||||
if !ok {
|
||||
warnings = append(warnings, "未知资产类型，已跳过: "+at)
|
||||
continue
|
||||
}
|
||||
items, truncated, ferr := s.ssPageAll(ctx, cl, ssType, filter, pageSize)
|
||||
if ferr != nil {
|
||||
errs = append(errs, fmt.Sprintf("%s(%s) 拉取失败: %v", at, target, ferr))
|
||||
continue
|
||||
}
|
||||
if truncated {
|
||||
warnings = append(warnings, fmt.Sprintf("%s(%s) 达到 %d 条上限，已截断", at, target, syncMaxPerType))
|
||||
}
|
||||
for _, raw := range items {
|
||||
if e := s.ssIngest(as, at, raw, synced); e != "" {
|
||||
errs = append(errs, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rebuild company attribution so freshly-synced assets attach to their company
|
||||
// (scope was written above; assets came after).
|
||||
if madeCompany {
|
||||
_ = cs.RecomputeAttribution()
|
||||
}
|
||||
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"synced": synced,
|
||||
"companies": companies,
|
||||
"warnings": warnings,
|
||||
"errors": errs,
|
||||
})
|
||||
}
|
||||
|
||||
// ssProjectMeta fetches a project's display name and its target root domains
|
||||
// (get_project.target is a newline-separated list) for company + scope creation.
|
||||
func (s *Server) ssProjectMeta(ctx context.Context, cl *mcphttp.Client, projectID string) (name string, roots []string, err error) {
|
||||
text, err := cl.Call(ctx, "get_project", map[string]any{"id": projectID})
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
var p struct {
|
||||
Name string `json:"name"`
|
||||
Target string `json:"target"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &p); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
for _, line := range strings.Split(p.Target, "\n") {
|
||||
if t := strings.TrimSpace(line); t != "" {
|
||||
roots = append(roots, t)
|
||||
}
|
||||
}
|
||||
return p.Name, roots, nil
|
||||
}
|
||||
|
||||
// ssPageAll pages through list_assets for one asset_type + filter until a short/empty
|
||||
// page or the protection cap. The {list:[...]} envelope has no total, so we stop when
|
||||
// a page returns fewer than pageSize items.
|
||||
func (s *Server) ssPageAll(ctx context.Context, cl *mcphttp.Client, ssType string, filter map[string]any, pageSize int) (items []json.RawMessage, truncated bool, err error) {
|
||||
for page := 1; ; page++ {
|
||||
args := map[string]any{"asset_type": ssType, "pageIndex": page, "pageSize": pageSize}
|
||||
if len(filter) > 0 {
|
||||
args["filter"] = filter
|
||||
}
|
||||
text, cerr := cl.Call(ctx, "list_assets", args)
|
||||
if cerr != nil {
|
||||
return items, truncated, cerr
|
||||
}
|
||||
var env struct {
|
||||
List []json.RawMessage `json:"list"`
|
||||
}
|
||||
if uerr := json.Unmarshal([]byte(text), &env); uerr != nil {
|
||||
return items, truncated, uerr
|
||||
}
|
||||
if len(env.List) == 0 {
|
||||
break
|
||||
}
|
||||
items = append(items, env.List...)
|
||||
if len(items) >= syncMaxPerType {
|
||||
items = items[:syncMaxPerType]
|
||||
truncated = true
|
||||
break
|
||||
}
|
||||
if len(env.List) < pageSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
return items, truncated, nil
|
||||
}
|
||||
|
||||
// ssIngest maps one ScopeSentry asset JSON to the ARTEX asset store and upserts it.
|
||||
// Returns a non-empty error string on failure. synced is incremented per kind.
|
||||
func (s *Server) ssIngest(as *db.AssetStore, assetType string, raw json.RawMessage, synced map[string]int) string {
|
||||
switch assetType {
|
||||
case "subdomain":
|
||||
var it struct {
|
||||
Host string `json:"host"`
|
||||
Type string `json:"type"`
|
||||
Value []string `json:"value"`
|
||||
IP []string `json:"ip"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "subdomain 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Host == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertSubdomain(db.UpsertSubdomainReq{Domain: it.Host, RecordType: it.Type, RecordValue: it.Value}); err != nil {
|
||||
return "subdomain " + it.Host + ": " + err.Error()
|
||||
}
|
||||
synced["subdomain"]++
|
||||
for _, ip := range it.IP {
|
||||
if ip != "" {
|
||||
if _, err := as.UpsertIP(db.UpsertIPReq{IP: ip, BoundDomains: []string{it.Host}}); err == nil {
|
||||
synced["ip"]++
|
||||
}
|
||||
}
|
||||
}
|
||||
case "app":
|
||||
var it struct {
|
||||
Name string `json:"name"`
|
||||
Category string `json:"category"`
|
||||
Description string `json:"description"`
|
||||
ICP string `json:"icp"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "app 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Name == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertApp(db.UpsertAppReq{Name: it.Name, Category: it.Category, Description: it.Description, ICP: it.ICP}); err != nil {
|
||||
return "app " + it.Name + ": " + err.Error()
|
||||
}
|
||||
synced["app"]++
|
||||
case "service":
|
||||
var it struct {
|
||||
Domain string `json:"domain"`
|
||||
IP string `json:"ip"`
|
||||
Port string `json:"port"`
|
||||
Service string `json:"service"`
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
Status *int `json:"status"`
|
||||
Products []string `json:"products"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "service 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Service == "http" || it.URL != "" {
|
||||
if it.URL == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertHTTPService(db.UpsertHTTPServiceReq{
|
||||
URL: it.URL, Technologies: it.Products, StatusCode: it.Status,
|
||||
PageTitle: it.Title, FaviconMMH3: it.Icon, IP: it.IP,
|
||||
}); err != nil {
|
||||
return "service " + it.URL + ": " + err.Error()
|
||||
}
|
||||
} else {
|
||||
port, _ := strconv.Atoi(it.Port)
|
||||
if _, err := as.UpsertOtherService(db.UpsertOtherServiceReq{
|
||||
Domain: it.Domain, IP: it.IP, Port: port, ServiceName: it.Service,
|
||||
}); err != nil {
|
||||
return "service " + it.IP + ":" + it.Port + ": " + err.Error()
|
||||
}
|
||||
}
|
||||
synced["service"]++
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// queryInt reads an int query param with a default fallback.
|
||||
func queryInt(r *http.Request, key string, def int) int {
|
||||
if v := r.URL.Query().Get(key); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return def
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
_ "github.com/jackc/pgx/v5/stdlib"
|
||||
)
|
||||
|
||||
// TestMain acquires a PostgreSQL advisory lock (7337741002) for the entire
|
||||
// server test suite so cross-package DELETE cleanup races with db/agent
|
||||
// packages are avoided when running `go test ./...`.
|
||||
func TestMain(m *testing.M) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
conn, err := sql.Open("pgx", dsn)
|
||||
if err != nil || conn.Ping() != nil {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, err := conn.Exec(`SELECT pg_advisory_lock(7337741002)`); err != nil {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
defer conn.Exec(`SELECT pg_advisory_unlock(7337741002)`) //nolint:errcheck
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
func names(tools []actool.CoreTool) map[string]actool.CoreTool {
|
||||
m := map[string]actool.CoreTool{}
|
||||
for _, t := range tools {
|
||||
m[t.Name()] = t
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// TestWireTools verifies end-to-end against the live dev PG: wireTools seeds the
|
||||
// catalog, ToolResolve keeps record_fact for worker but drops it for planner (not
|
||||
// bound), and an edited description + injected default flow through.
|
||||
func TestWireTools(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v) — skipping", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
|
||||
wireTools(pg, nil) // nil domainReg: test only covers filter/decoration, not injection
|
||||
t.Cleanup(func() { agent.ToolResolve = nil })
|
||||
// Prevent interactive-shell Bash decoration: if the worker agent has
|
||||
// interactive_shell=true in the DB, ToolResolve would append a note to
|
||||
// Bash's description. Disable that for this test so the undecorated
|
||||
// pass-through assertion holds regardless of DB state.
|
||||
t.Setenv("AGENT_CORE_DISABLE_INTERACTIVE_SHELL", "1")
|
||||
|
||||
// Seeding populated the catalog.
|
||||
rf, err := pg.GetTool("record_fact")
|
||||
if err != nil || rf == nil {
|
||||
t.Fatalf("record_fact not seeded: %v", err)
|
||||
}
|
||||
|
||||
// Build a base as the worker does: its domain tools + defaults, resolved for
|
||||
// "worker" then for "planner".
|
||||
ts := agent.NewToolSet(nil, "")
|
||||
base := append(ts.WorkerTools(), actool.DefaultTools()...)
|
||||
ctx := context.Background()
|
||||
|
||||
worker := names(agent.ToolResolve(ctx, "worker", base))
|
||||
if _, ok := worker["record_fact"]; !ok {
|
||||
t.Error("worker lost record_fact")
|
||||
}
|
||||
// SDK generic tools aren't seeded → ToolResolve passes them through unchanged
|
||||
// (same object, original description — not decorated). No startup prune touches
|
||||
// non-catalog rows, so future user-defined custom tools survive too.
|
||||
if bash, ok := worker["Bash"]; !ok {
|
||||
t.Error("worker lost Bash (should pass through)")
|
||||
} else if bash.Description() != actool.NewBash().Description() {
|
||||
t.Error("Bash should pass through undecorated, but description changed")
|
||||
}
|
||||
// planner is not bound to record_fact → resolving a base that contains it drops it.
|
||||
planner := names(agent.ToolResolve(ctx, "planner", base))
|
||||
if _, ok := planner["record_fact"]; ok {
|
||||
t.Error("planner should not get record_fact (not bound)")
|
||||
}
|
||||
|
||||
// Edit description + add a default to a scalar param, then confirm the resolved
|
||||
// worker tool reflects both. Restore from code default on cleanup.
|
||||
t.Cleanup(func() { _ = pg.UpsertToolForce(rf.Key, rf.Description, rf.Schema, mustJSON(rf.Agents)) })
|
||||
var schema map[string]any
|
||||
_ = json.Unmarshal(rf.Schema, &schema)
|
||||
if props, ok := schema["properties"].(map[string]any); ok {
|
||||
if conf, ok := props["confidence"].(map[string]any); ok {
|
||||
conf["default"] = "inferred"
|
||||
}
|
||||
}
|
||||
edited := mustJSON(schema)
|
||||
if err := pg.UpdateTool(rf.Key, "EDITED DESC", edited, mustJSON([]string{"worker"}), true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
worker2 := names(agent.ToolResolve(ctx, "worker", base))
|
||||
got := worker2["record_fact"]
|
||||
if got == nil {
|
||||
t.Fatal("record_fact missing after edit")
|
||||
}
|
||||
if got.Description() != "EDITED DESC" {
|
||||
t.Errorf("description = %q, want EDITED DESC", got.Description())
|
||||
}
|
||||
// Default injection: omit confidence → handler input should carry the default.
|
||||
// We can't run the real handler (nil store), but InputSchema must show the default.
|
||||
sc := got.InputSchema()
|
||||
props := sc["properties"].(map[string]any)
|
||||
conf := props["confidence"].(map[string]any)
|
||||
if conf["default"] != "inferred" {
|
||||
t.Errorf("confidence.default = %v, want inferred", conf["default"])
|
||||
}
|
||||
}
|
||||
|
||||
func mustJSON(v any) json.RawMessage {
|
||||
b, _ := json.Marshal(v)
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// ---------- P3 agent triggers (仅自定义 agent) ----------
|
||||
|
||||
func (s *Server) pgListTriggers(w http.ResponseWriter, r *http.Request) {
|
||||
pg, a, ok := s.agentByKey(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
trs, err := pg.ListTriggersFor(a.Key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"triggers": trs})
|
||||
}
|
||||
|
||||
type triggerReq struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
IntervalSec int `json:"interval_sec"`
|
||||
OnFinding bool `json:"on_finding"`
|
||||
OnGoalMet bool `json:"on_goal_met"`
|
||||
OnTaskTimeout bool `json:"on_task_timeout"`
|
||||
IntervalMessage string `json:"interval_message"`
|
||||
FindingMessage string `json:"finding_message"`
|
||||
GoalMessage string `json:"goal_message"`
|
||||
TaskTimeoutMessage string `json:"task_timeout_message"`
|
||||
}
|
||||
|
||||
func (s *Server) pgCreateTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg, a, ok := s.agentByKey(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if a.Builtin {
|
||||
writeErr(w, 400, "触发器仅支持自定义 agent")
|
||||
return
|
||||
}
|
||||
var req triggerReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.IntervalSec < 0 {
|
||||
req.IntervalSec = 0
|
||||
}
|
||||
if req.IntervalSec == 0 && !req.OnFinding && !req.OnGoalMet && !req.OnTaskTimeout {
|
||||
writeErr(w, 400, "至少选择一种触发条件(定时/发现finding/目标达成/任务超时)")
|
||||
return
|
||||
}
|
||||
tr, err := pg.CreateTrigger(&db.AgentTrigger{
|
||||
AgentKey: a.Key, Enabled: req.Enabled, IntervalSec: req.IntervalSec,
|
||||
OnFinding: req.OnFinding, OnGoalMet: req.OnGoalMet, OnTaskTimeout: req.OnTaskTimeout,
|
||||
IntervalMessage: req.IntervalMessage, FindingMessage: req.FindingMessage,
|
||||
GoalMessage: req.GoalMessage, TaskTimeoutMessage: req.TaskTimeoutMessage,
|
||||
})
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, tr)
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad trigger id")
|
||||
return
|
||||
}
|
||||
var req triggerReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.IntervalSec < 0 {
|
||||
req.IntervalSec = 0
|
||||
}
|
||||
if req.IntervalSec == 0 && !req.OnFinding && !req.OnGoalMet && !req.OnTaskTimeout {
|
||||
writeErr(w, 400, "至少选择一种触发条件(定时/发现finding/目标达成/任务超时)")
|
||||
return
|
||||
}
|
||||
if err := pg.UpdateTrigger(&db.AgentTrigger{
|
||||
ID: id, Enabled: req.Enabled, IntervalSec: req.IntervalSec,
|
||||
OnFinding: req.OnFinding, OnGoalMet: req.OnGoalMet, OnTaskTimeout: req.OnTaskTimeout,
|
||||
IntervalMessage: req.IntervalMessage, FindingMessage: req.FindingMessage,
|
||||
GoalMessage: req.GoalMessage, TaskTimeoutMessage: req.TaskTimeoutMessage,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad trigger id")
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteTrigger(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": id})
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
//go:build embedui
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// webuiDist holds the statically-exported frontend (built with
|
||||
// `cd web && npm run build:static`, then copied into server/webui/dist). Compiled
|
||||
// into the binary only when building with `-tags embedui`. The `all:` prefix is
|
||||
// required so Next's `_next/` asset dir (leading underscore) is included.
|
||||
//
|
||||
//go:embed all:webui/dist
|
||||
var webuiDist embed.FS
|
||||
|
||||
// webuiHandler serves the embedded SPA. Public (no JWT) — auth is enforced
|
||||
// client-side and on /api. Serves the exported per-route index.html files and
|
||||
// falls back to index.html so client-side routing still resolves unknown paths.
|
||||
func (s *Server) webuiHandler() http.Handler {
|
||||
root, err := fs.Sub(webuiDist, "webui/dist")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
p := strings.TrimPrefix(path.Clean(r.URL.Path), "/")
|
||||
if p == "" {
|
||||
p = "index.html"
|
||||
}
|
||||
// 1) exact file (assets: _next/*, favicon.ico, ...)
|
||||
// 2) route dir → <p>/index.html (trailingSlash export)
|
||||
// 3) <p>.html
|
||||
// 4) SPA 兜底 → index.html（交给客户端路由）
|
||||
for _, cand := range []string{p, p + "/index.html", p + ".html"} {
|
||||
if serveFileIfExists(w, r, root, cand) {
|
||||
return
|
||||
}
|
||||
}
|
||||
http.ServeFileFS(w, r, root, "index.html")
|
||||
})
|
||||
}
|
||||
|
||||
// serveFileIfExists serves name from fsys when it exists as a regular file.
|
||||
func serveFileIfExists(w http.ResponseWriter, r *http.Request, fsys fs.FS, name string) bool {
|
||||
f, err := fsys.Open(name)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
st, statErr := f.Stat()
|
||||
_ = f.Close()
|
||||
if statErr != nil || st.IsDir() {
|
||||
return false
|
||||
}
|
||||
http.ServeFileFS(w, r, fsys, name)
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
//go:build !embedui
|
||||
|
||||
package server
|
||||
|
||||
import "net/http"
|
||||
|
||||
// webuiHandler is the no-embed stub (default build). The frontend is NOT bundled
|
||||
// into this binary — run it separately with `cd web && npm run dev` during
|
||||
// development. Build the bundled single-binary with:
|
||||
//
|
||||
// cd web && npm run build:static # produces web/out
|
||||
// cp -r web/out server/webui/dist # (or use the build script)
|
||||
// go build -tags embedui ./cmd/artex
|
||||
func (s *Server) webuiHandler() http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
http.Error(w, "前端未内嵌到此二进制（开发用 next dev；发布用 -tags embedui 构建）", http.StatusNotFound)
|
||||
})
|
||||
}
|
||||
Reference in new issue
Block a user