mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
Autumn-27andClaude Opus 4.8
0ef8561411
ARTEX — AI 自主渗透测试系统（Go 后端 + Next.js 前端，单二进制内嵌）。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
685 lines
25 KiB
Go
685 lines
25 KiB
Go
package db
|
||
|
||
import (
|
||
"database/sql"
|
||
"encoding/json"
|
||
"time"
|
||
)
|
||
|
||
// ---------- LLM profiles ----------
|
||
|
||
type LLMProfile struct {
|
||
ID int64 `json:"id"`
|
||
Name string `json:"name"`
|
||
Format string `json:"format"`
|
||
BaseURL string `json:"base_url,omitempty"`
|
||
Proxy string `json:"proxy,omitempty"` // LLM 出站代理(http/https/socks5);空=用环境变量
|
||
Model string `json:"model"`
|
||
APIKey string `json:"-"` // never serialized to UI
|
||
APIKeyHint string `json:"api_key_hint,omitempty"`
|
||
RatePerSecond float64 `json:"rate_per_second"`
|
||
RatePerMinute float64 `json:"rate_per_minute"`
|
||
// ContextWindowK is the model's context window in K tokens, used to size
|
||
// compaction thresholds. 0 = use a 200K default; capped at 1000 (1M).
|
||
ContextWindowK int `json:"context_window_k"`
|
||
// ReasoningEffort is the thinking-mode selector. "" = 默认(不发送思考参数);
|
||
// "off" = 显式关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" =
|
||
// 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider.
|
||
ReasoningEffort string `json:"reasoning_effort"`
|
||
IsDefault bool `json:"is_default"`
|
||
}
|
||
|
||
func (d *DB) ListProfiles() ([]*LLMProfile, error) {
|
||
rows, err := d.Query(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles ORDER BY id`)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []*LLMProfile
|
||
for rows.Next() {
|
||
var p LLMProfile
|
||
if err := rows.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKeyHint, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, &p)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// ActiveProfile returns the default (active) profile with its api key, or nil.
|
||
func (d *DB) ActiveProfile() (*LLMProfile, error) {
|
||
var p LLMProfile
|
||
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE is_default LIMIT 1`).
|
||
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
|
||
if err == sql.ErrNoRows {
|
||
return nil, nil
|
||
}
|
||
return &p, err
|
||
}
|
||
|
||
// ProfileByID returns one profile with its api key by id, or nil if not found.
|
||
// Used to run a task on a specific (non-default) LLM profile.
|
||
func (d *DB) ProfileByID(id int64) (*LLMProfile, error) {
|
||
var p LLMProfile
|
||
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE id=$1`, id).
|
||
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
|
||
if err == sql.ErrNoRows {
|
||
return nil, nil
|
||
}
|
||
return &p, err
|
||
}
|
||
|
||
// SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
|
||
func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
|
||
hint := p.APIKeyHint
|
||
if len(p.APIKey) >= 4 {
|
||
hint = "…" + p.APIKey[len(p.APIKey)-4:]
|
||
}
|
||
if p.ID == 0 {
|
||
var id int64
|
||
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort)
|
||
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11) RETURNING id`,
|
||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort).Scan(&id)
|
||
return id, err
|
||
}
|
||
if p.APIKey == "" {
|
||
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9 WHERE id=$10`,
|
||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
|
||
return p.ID, err
|
||
}
|
||
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11 WHERE id=$12`,
|
||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
|
||
return p.ID, err
|
||
}
|
||
|
||
func (d *DB) DeleteProfile(id int64) error {
|
||
_, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id)
|
||
return err
|
||
}
|
||
|
||
// SetActiveProfile makes one profile the global default (single-default invariant).
|
||
func (d *DB) SetActiveProfile(id int64) error {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil {
|
||
return err
|
||
}
|
||
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// ---------- Agents / prompts ----------
|
||
|
||
type Agent struct {
|
||
ID int64 `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"` // 单次运行最大轮次;0=不限制
|
||
RunSecs int `json:"run_seconds"` // worker 单次运行墙钟上限(秒);0=不限制
|
||
WebSearch bool `json:"web_search"` // 是否启用网络搜索(受系统全局开关门控)
|
||
InteractiveShell bool `json:"interactive_shell"` // 是否启用交互式 shell(持久 PTY 会话工具族)
|
||
WrapupPrompt string `json:"wrapup_prompt"` // 收尾提示词(超时/步数耗尽时的 settlement 提示);空=用代码内置默认
|
||
WrapupMaxTurns int `json:"wrapup_max_turns"` // 收尾阶段自身的轮数预算;0=用代码内置默认(按 agent)
|
||
// 任务级超时收尾词(与 per-run 两套;仅 worker/planner 用);空/0=用代码内置默认。
|
||
TaskTimeoutWrapupPrompt string `json:"task_timeout_wrapup_prompt"`
|
||
TaskTimeoutWrapupMaxTurns int `json:"task_timeout_wrapup_max_turns"`
|
||
}
|
||
|
||
const agentCols = `id,key,name,COALESCE(description,''),role,builtin,enabled,COALESCE(max_turns,0),COALESCE(run_seconds,600),COALESCE(web_search,false),COALESCE(interactive_shell,false),COALESCE(wrapup_prompt,''),COALESCE(wrapup_max_turns,0),COALESCE(task_timeout_wrapup_prompt,''),COALESCE(task_timeout_wrapup_max_turns,0)`
|
||
|
||
func scanAgent(sc interface{ Scan(...any) error }) (*Agent, error) {
|
||
var a Agent
|
||
err := sc.Scan(&a.ID, &a.Key, &a.Name, &a.Description, &a.Role, &a.Builtin, &a.Enabled, &a.MaxTurns, &a.RunSecs, &a.WebSearch, &a.InteractiveShell, &a.WrapupPrompt, &a.WrapupMaxTurns, &a.TaskTimeoutWrapupPrompt, &a.TaskTimeoutWrapupMaxTurns)
|
||
return &a, err
|
||
}
|
||
|
||
func (d *DB) ListAgents() ([]*Agent, error) {
|
||
rows, err := d.Query(`SELECT ` + agentCols + ` FROM agents ORDER BY id`)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []*Agent
|
||
for rows.Next() {
|
||
a, err := scanAgent(rows)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, a)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
func (d *DB) GetAgentByKey(key string) (*Agent, error) {
|
||
a, err := scanAgent(d.QueryRow(`SELECT `+agentCols+` FROM agents WHERE key=$1`, key))
|
||
if err == sql.ErrNoRows {
|
||
return nil, nil
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return a, nil
|
||
}
|
||
|
||
// AgentBindingCounts returns per-agent binding counts in a few grouped queries
|
||
// (NO N+1): visible MCP servers and skills keyed by agent id, and bound tools
|
||
// keyed by agent key (tools.agents is a JSONB array of agent keys). Missing keys
|
||
// mean zero. Used to show "MCP N · Skill N · 工具 N" on the agent cards.
|
||
func (d *DB) AgentBindingCounts() (mcp map[int64]int, skill map[int64]int, tools map[string]int, err error) {
|
||
mcp, skill, tools = map[int64]int{}, map[int64]int{}, map[string]int{}
|
||
byID := func(q string, into map[int64]int) error {
|
||
rows, e := d.Query(q)
|
||
if e != nil {
|
||
return e
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var id int64
|
||
var n int
|
||
if e := rows.Scan(&id, &n); e != nil {
|
||
return e
|
||
}
|
||
into[id] = n
|
||
}
|
||
return rows.Err()
|
||
}
|
||
if err = byID(`SELECT agent_id, count(DISTINCT resource_id) FROM agent_visibility WHERE resource_kind='mcp' AND enabled GROUP BY agent_id`, mcp); err != nil {
|
||
return
|
||
}
|
||
if err = byID(`SELECT agent_id, count(*) FROM agent_skill_visibility WHERE enabled GROUP BY agent_id`, skill); err != nil {
|
||
return
|
||
}
|
||
rows, e := d.Query(`SELECT elem, count(*) FROM tools, jsonb_array_elements_text(agents) AS elem GROUP BY elem`)
|
||
if e != nil {
|
||
err = e
|
||
return
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var k string
|
||
var n int
|
||
if e := rows.Scan(&k, &n); e != nil {
|
||
err = e
|
||
return
|
||
}
|
||
tools[k] = n
|
||
}
|
||
err = rows.Err()
|
||
return
|
||
}
|
||
|
||
// CreateAgent inserts a custom (builtin=false) conversational agent with role
|
||
// 'assistant'. Returns the new row. Callers validate key/name upstream; the DB
|
||
// enforces key charset + uniqueness and the role check constraint.
|
||
func (d *DB) CreateAgent(key, name, description string) (*Agent, error) {
|
||
a, err := scanAgent(d.QueryRow(`
|
||
INSERT INTO agents(key, name, description, role, builtin, enabled)
|
||
VALUES ($1, $2, NULLIF($3,''), 'assistant', false, true)
|
||
RETURNING `+agentCols, key, name, description))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return a, nil
|
||
}
|
||
|
||
// UpdateAgentMeta updates a custom agent's display name + description. Built-in
|
||
// agents are left untouched (guarded by the caller / the builtin flag).
|
||
func (d *DB) UpdateAgentMeta(key, name, description string) error {
|
||
_, err := d.Exec(`UPDATE agents SET name=$2, description=NULLIF($3,'') WHERE key=$1 AND builtin=false`, key, name, description)
|
||
return err
|
||
}
|
||
|
||
// DeleteAgent removes a custom agent. Built-in agents are protected by the
|
||
// builtin=false guard. agent_prompts / agent_prompt_vars / visibility rows cascade
|
||
// via FK; tools.agents bindings for the key are cleaned by the caller.
|
||
func (d *DB) DeleteAgent(key string) error {
|
||
_, err := d.Exec(`DELETE FROM agents WHERE key=$1 AND builtin=false`, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentMaxTurns updates an agent's max_turns (0 = unlimited).
|
||
func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error {
|
||
if maxTurns < 0 {
|
||
maxTurns = 0
|
||
}
|
||
_, err := d.Exec(`UPDATE agents SET max_turns=$1 WHERE key=$2`, maxTurns, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentWebSearch toggles whether an agent uses network search (still gated by
|
||
// the global web-search master switch + backend/key config).
|
||
func (d *DB) SetAgentWebSearch(key string, on bool) error {
|
||
_, err := d.Exec(`UPDATE agents SET web_search=$1 WHERE key=$2`, on, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentInteractiveShell toggles whether an agent gets the interactive shell
|
||
// (持久 PTY 会话) tool family + Bash 提示词联动(见 docs/交互式shell设计.md §14.2).
|
||
func (d *DB) SetAgentInteractiveShell(key string, on bool) error {
|
||
_, err := d.Exec(`UPDATE agents SET interactive_shell=$1 WHERE key=$2`, on, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentWrapupPrompt stores an agent's wrap-up (settlement) prompt. Empty string
|
||
// means "use the code built-in default" — resolved at runtime by resolveWrapup.
|
||
func (d *DB) SetAgentWrapupPrompt(key, prompt string) error {
|
||
_, err := d.Exec(`UPDATE agents SET wrapup_prompt=$1 WHERE key=$2`, prompt, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentWrapupMaxTurns stores the wrap-up phase's own turn budget. 0 means "use
|
||
// the code built-in default" — resolved at runtime by resolveWrapupTurns.
|
||
func (d *DB) SetAgentWrapupMaxTurns(key string, n int) error {
|
||
_, err := d.Exec(`UPDATE agents SET wrapup_max_turns=$1 WHERE key=$2`, n, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentTaskTimeoutWrapup stores an agent's task-timeout wrap-up prompt (empty =
|
||
// use code built-in default; only worker/planner have one) and its turn budget
|
||
// (0 = default). Resolved at runtime by resolveTaskTimeoutWrapup / …Turns.
|
||
func (d *DB) SetAgentTaskTimeoutWrapup(key, prompt string, maxTurns int) error {
|
||
_, err := d.Exec(`UPDATE agents SET task_timeout_wrapup_prompt=$1, task_timeout_wrapup_max_turns=$2 WHERE key=$3`, prompt, maxTurns, key)
|
||
return err
|
||
}
|
||
|
||
// SetAgentRunSeconds updates an agent's run_seconds wall-clock budget (0 = unlimited).
|
||
func (d *DB) SetAgentRunSeconds(key string, runSecs int) error {
|
||
if runSecs < 0 {
|
||
runSecs = 0
|
||
}
|
||
_, err := d.Exec(`UPDATE agents SET run_seconds=$1 WHERE key=$2`, runSecs, key)
|
||
return err
|
||
}
|
||
|
||
type PromptVar struct {
|
||
Name string `json:"name"`
|
||
Description string `json:"description"`
|
||
Example string `json:"example"`
|
||
Source string `json:"source"`
|
||
}
|
||
|
||
func (d *DB) PromptVars(agentID int64) ([]PromptVar, error) {
|
||
rows, err := d.Query(`SELECT var_name,COALESCE(description,''),COALESCE(example,''),source FROM agent_prompt_vars WHERE agent_id=$1 ORDER BY var_name`, agentID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []PromptVar{}
|
||
for rows.Next() {
|
||
var v PromptVar
|
||
if err := rows.Scan(&v.Name, &v.Description, &v.Example, &v.Source); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, v)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// CurrentPrompt returns the agent's active template text ("" if none set yet).
|
||
func (d *DB) CurrentPrompt(agentID int64) (string, error) {
|
||
var tmpl sql.NullString
|
||
err := d.QueryRow(`SELECT p.template_text FROM agents a JOIN agent_prompts p ON p.id=a.current_prompt_id WHERE a.id=$1`, agentID).Scan(&tmpl)
|
||
if err == sql.ErrNoRows {
|
||
return "", nil
|
||
}
|
||
return tmpl.String, err
|
||
}
|
||
|
||
// SeedPromptIfEmpty writes the code-default template as the agent's first prompt
|
||
// version ONLY when it has none yet (current_prompt_id IS NULL). Mirrors
|
||
// SeedTool's first-insert-only philosophy: a user's edited prompt is never
|
||
// clobbered on restart. Idempotent — a no-op once any version exists.
|
||
func (d *DB) SeedPromptIfEmpty(agentID int64, tmpl string) error {
|
||
var cur sql.NullInt64
|
||
if err := d.QueryRow(`SELECT current_prompt_id FROM agents WHERE id=$1`, agentID).Scan(&cur); err != nil {
|
||
return err
|
||
}
|
||
if cur.Valid {
|
||
return nil // already seeded or user-edited → leave it
|
||
}
|
||
_, err := d.SavePrompt(agentID, tmpl, "内置默认", "system")
|
||
return err
|
||
}
|
||
|
||
// ResetPromptToDefault appends the code-default template as a new version and
|
||
// points current at it — the explicit "恢复为内置默认" action.
|
||
func (d *DB) ResetPromptToDefault(agentID int64, tmpl string) (int, error) {
|
||
return d.SavePrompt(agentID, tmpl, "恢复为内置默认", "system")
|
||
}
|
||
|
||
// SavePrompt appends a new version and points current_prompt_id at it.
|
||
func (d *DB) SavePrompt(agentID int64, template, note, by string) (int, error) {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
defer tx.Rollback()
|
||
var ver int
|
||
if err := tx.QueryRow(`SELECT COALESCE(max(version),0)+1 FROM agent_prompts WHERE agent_id=$1`, agentID).Scan(&ver); err != nil {
|
||
return 0, err
|
||
}
|
||
var pid int64
|
||
if err := tx.QueryRow(`INSERT INTO agent_prompts(agent_id,version,template_text,note,updated_by) VALUES ($1,$2,$3,NULLIF($4,''),NULLIF($5,'')) RETURNING id`,
|
||
agentID, ver, template, note, by).Scan(&pid); err != nil {
|
||
return 0, err
|
||
}
|
||
if _, err := tx.Exec(`UPDATE agents SET current_prompt_id=$1 WHERE id=$2`, pid, agentID); err != nil {
|
||
return 0, err
|
||
}
|
||
return ver, tx.Commit()
|
||
}
|
||
|
||
type PromptVersion struct {
|
||
Version int `json:"version"`
|
||
Template string `json:"template_text"`
|
||
Note string `json:"note,omitempty"`
|
||
CreatedAt time.Time `json:"ts"`
|
||
}
|
||
|
||
func (d *DB) ListPromptVersions(agentID int64) ([]PromptVersion, error) {
|
||
rows, err := d.Query(`SELECT version,template_text,COALESCE(note,''),created_at FROM agent_prompts WHERE agent_id=$1 ORDER BY version DESC`, agentID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []PromptVersion{}
|
||
for rows.Next() {
|
||
var v PromptVersion
|
||
if err := rows.Scan(&v.Version, &v.Template, &v.Note, &v.CreatedAt); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, v)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// ---------- MCP servers ----------
|
||
|
||
type MCPServer struct {
|
||
ID int64 `json:"id"`
|
||
Name string `json:"name"`
|
||
Transport string `json:"transport"`
|
||
Command string `json:"command,omitempty"`
|
||
Args json.RawMessage `json:"args"`
|
||
Env json.RawMessage `json:"env"`
|
||
URL string `json:"url,omitempty"`
|
||
Enabled bool `json:"enabled"`
|
||
Tools []string `json:"tools,omitempty"` // cached tool names (mcp_tools_cache)
|
||
}
|
||
|
||
func (d *DB) ListMCP() ([]*MCPServer, error) {
|
||
rows, err := d.Query(`SELECT id,name,transport,COALESCE(command,''),args,env,COALESCE(url,''),enabled FROM mcp_servers ORDER BY id`)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
var out []*MCPServer
|
||
for rows.Next() {
|
||
var m MCPServer
|
||
var args, env []byte
|
||
if err := rows.Scan(&m.ID, &m.Name, &m.Transport, &m.Command, &args, &env, &m.URL, &m.Enabled); err != nil {
|
||
rows.Close()
|
||
return nil, err
|
||
}
|
||
m.Args, m.Env = json.RawMessage(args), json.RawMessage(env)
|
||
out = append(out, &m)
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
rows.Close()
|
||
return nil, err
|
||
}
|
||
rows.Close() // free the connection before the per-server tool-cache queries below
|
||
// Attach each server's cached tool names (best-effort; empty until discovered).
|
||
for _, m := range out {
|
||
m.Tools, _ = d.MCPToolNames(m.ID)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// MCPToolNames returns the cached tool names for a server (empty until discovered).
|
||
func (d *DB) MCPToolNames(serverID int64) ([]string, error) {
|
||
rows, err := d.Query(`SELECT tool_name FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []string
|
||
for rows.Next() {
|
||
var n string
|
||
if err := rows.Scan(&n); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, n)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// MCPTool is one cached tool of an MCP server (name + description).
|
||
type MCPTool struct {
|
||
Name string `json:"name"`
|
||
Description string `json:"description"`
|
||
}
|
||
|
||
// MCPToolsDetailed returns the cached tools (name + description) for a server.
|
||
func (d *DB) MCPToolsDetailed(serverID int64) ([]MCPTool, error) {
|
||
rows, err := d.Query(`SELECT tool_name, COALESCE(description,'') FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []MCPTool
|
||
for rows.Next() {
|
||
var t MCPTool
|
||
if err := rows.Scan(&t.Name, &t.Description); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, t)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SaveMCPTools replaces the cached tool list for a server (called after discovery).
|
||
func (d *DB) SaveMCPTools(serverID int64, tools []MCPTool) error {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
if _, err := tx.Exec(`DELETE FROM mcp_tools_cache WHERE server_id=$1`, serverID); err != nil {
|
||
return err
|
||
}
|
||
for _, t := range tools {
|
||
if _, err := tx.Exec(`INSERT INTO mcp_tools_cache(server_id, tool_name, description) VALUES ($1,$2,$3)
|
||
ON CONFLICT (server_id, tool_name) DO UPDATE SET description=EXCLUDED.description`, serverID, t.Name, t.Description); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
func (d *DB) SaveMCP(m *MCPServer) (int64, error) {
|
||
args, env := string(m.Args), string(m.Env)
|
||
if args == "" {
|
||
args = "[]"
|
||
}
|
||
if env == "" {
|
||
env = "{}"
|
||
}
|
||
if m.ID == 0 {
|
||
var id int64
|
||
err := d.QueryRow(`INSERT INTO mcp_servers(name,transport,command,args,env,url,enabled) VALUES ($1,$2,NULLIF($3,''),$4,$5,NULLIF($6,''),$7) RETURNING id`,
|
||
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled).Scan(&id)
|
||
return id, err
|
||
}
|
||
_, err := d.Exec(`UPDATE mcp_servers SET name=$1,transport=$2,command=NULLIF($3,''),args=$4,env=$5,url=NULLIF($6,''),enabled=$7 WHERE id=$8`,
|
||
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled, m.ID)
|
||
return m.ID, err
|
||
}
|
||
|
||
func (d *DB) DeleteMCP(id int64) error {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE resource_kind='mcp' AND resource_id=$1`, id); err != nil {
|
||
return err
|
||
}
|
||
if _, err := tx.Exec(`DELETE FROM mcp_servers WHERE id=$1`, id); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// ---------- Skill visibility (agent × skill_name) ----------
|
||
|
||
// AgentSkillNames returns the skill directory names visible to an agent.
|
||
func (d *DB) AgentSkillNames(agentID int64) ([]string, error) {
|
||
rows, err := d.Query(`SELECT skill_name FROM agent_skill_visibility WHERE agent_id=$1 AND enabled ORDER BY skill_name`, agentID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []string
|
||
for rows.Next() {
|
||
var name string
|
||
if err := rows.Scan(&name); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, name)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SkillAgents returns the agent IDs that can see a skill.
|
||
func (d *DB) SkillAgents(skillName string) ([]int64, error) {
|
||
rows, err := d.Query(`SELECT agent_id FROM agent_skill_visibility WHERE skill_name=$1 AND enabled`, skillName)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []int64
|
||
for rows.Next() {
|
||
var id int64
|
||
if err := rows.Scan(&id); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, id)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SetAgentSkillVisibility replaces all skill visibility for an agent.
|
||
func (d *DB) SetAgentSkillVisibility(agentID int64, names []string) error {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
if _, err := tx.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1`, agentID); err != nil {
|
||
return err
|
||
}
|
||
for _, name := range names {
|
||
if _, err := tx.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
|
||
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, name); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// ToggleSkillVisibility sets one (agent, skill_name) visibility on/off.
|
||
func (d *DB) ToggleSkillVisibility(agentID int64, skillName string, on bool) error {
|
||
if on {
|
||
_, err := d.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
|
||
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, skillName)
|
||
return err
|
||
}
|
||
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1 AND skill_name=$2`, agentID, skillName)
|
||
return err
|
||
}
|
||
|
||
// DeleteSkillVisibility removes all visibility rows for a skill (called on skill delete).
|
||
func (d *DB) DeleteSkillVisibility(skillName string) error {
|
||
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE skill_name=$1`, skillName)
|
||
return err
|
||
}
|
||
|
||
// ---------- Visibility (agent × mcp) ----------
|
||
|
||
// AgentVisible returns the resource ids of a kind visible to an agent.
|
||
func (d *DB) AgentVisible(agentID int64, kind string) ([]int64, error) {
|
||
rows, err := d.Query(`SELECT resource_id FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND enabled`, agentID, kind)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []int64
|
||
for rows.Next() {
|
||
var id int64
|
||
if err := rows.Scan(&id); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, id)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// ResourceAgents returns the agent ids that can see a resource.
|
||
func (d *DB) ResourceAgents(kind string, resourceID int64) ([]int64, error) {
|
||
rows, err := d.Query(`SELECT agent_id FROM agent_visibility WHERE resource_kind=$1 AND resource_id=$2 AND enabled`, kind, resourceID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
var out []int64
|
||
for rows.Next() {
|
||
var id int64
|
||
if err := rows.Scan(&id); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, id)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SetAgentVisibilityKind replaces the full set of visible resources of a kind for
|
||
// an agent (agent-side bulk write). Bidirectional with the resource-side view —
|
||
// both read/write the same agent_visibility rows.
|
||
func (d *DB) SetAgentVisibilityKind(agentID int64, kind string, resourceIDs []int64) error {
|
||
tx, err := d.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2`, agentID, kind); err != nil {
|
||
return err
|
||
}
|
||
for _, rid := range resourceIDs {
|
||
if _, err := tx.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
|
||
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, rid); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// ToggleVisibility sets one (agent, kind, resource) visibility on/off (idempotent).
|
||
func (d *DB) ToggleVisibility(agentID int64, kind string, resourceID int64, on bool) error {
|
||
if on {
|
||
_, err := d.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
|
||
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, resourceID)
|
||
return err
|
||
}
|
||
_, err := d.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND resource_id=$3`, agentID, kind, resourceID)
|
||
return err
|
||
}
|