First Commit

2026-10-09 08:38:16 +08:00
commit 0335d572de
756 changed files with 201663 additions and 0 deletions
@@ -0,0 +1,373 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"log"
+	"path/filepath"
+	"strings"
+
+	"github.com/Autumn-27/artex/agent"
+	"github.com/Autumn-27/artex/db"
+	"github.com/Autumn-27/artex/traffic"
+	"github.com/Autumn-27/norma/skill"
+	actool "github.com/Autumn-27/norma/tool"
+)
+
+// 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)
+			}
+			// Attribution for the usage ledger: ToolAugment only gets (ctx, agentKey),
+			// so the run's task/session ids ride in on the ctx (agent.RunInfo). Read it
+			// once here — this closure is rebuilt per run, so the captured value always
+			// belongs to this run.
+			extra = append(extra, meterSkillTool(reg.Tool(), pg, reg, a.Key, agent.RunInfoFrom(ctx)))
+		}
+
+		// host tools (traffic / orchestration / custom) — added to every agent's base;
+		// 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) {
+	agent.FindingTrafficBindingEnabled = func() bool { return pg.GetBool(settingAgentTrafficBinding, false) }
+	// 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(ctx 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
+		}
+		runInfo := agent.RunInfoFrom(ctx)
+		resolve := func(t actool.CoreTool, row *db.Tool) actool.CoreTool {
+			var schema map[string]any
+			if len(row.Schema) > 0 {
+				_ = json.Unmarshal(row.Schema, &schema)
+			}
+			return meterTool(agent.DecorateTool(t, row.Description, schema), pg, row.Key, agentKey, runInfo)
+		}
+		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
+			}
+			out = append(out, resolve(t, row))
+		}
+		// 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()
+				}
+				out = append(out, resolve(inst, row))
+			}
+		}
+		// 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).
+//
+// The nil ExplorationStore is deliberate — these instances are task-less by
+// construction — so every tool here must tolerate it. Asset/company tools do
+// (they only need the AssetStore); the exploration-graph tools refuse with a
+// clear message via ToolSet.needExploration. Binding one of them to a task-less
+// agent in the tools table is therefore a useless tool, not a crash.
+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
+}