diff --git a/agent/assembly.go b/agent/assembly.go index f4f8e9a..c0885d9 100644 --- a/agent/assembly.go +++ b/agent/assembly.go @@ -6,6 +6,7 @@ import ( "fmt" "log" "runtime/debug" + "sync" actool "github.com/Autumn-27/norma/tool" ) @@ -23,15 +24,36 @@ type DeferredInfo struct { UnlockSkill func(skillName string) } -// ToolAugment, if set, returns the EXTRA tools an agent should see beyond its -// built-in base set — the agent's visible skills (packed into one Skill meta-tool) +// ToolAugmentFunc is the ToolAugment hook. If set, it returns the EXTRA tools +// an agent should see beyond its built-in base set — the agent's visible skills (packed into one Skill meta-tool) // and visible MCP servers (expanded to mcp__server__tool). It also returns the // DeferredInfo describing how those MCP tools are deferred/gated. The server wires // it to the PG agent_visibility table. cleanup releases any spawned MCP clients. // // When nil, agents run with only their built-in tools — behavior is unchanged // until a user assigns a skill/MCP to the agent in the UI. -var ToolAugment func(ctx context.Context, agentKey string) (extra []actool.CoreTool, def DeferredInfo, cleanup func()) +type ToolAugmentFunc func(ctx context.Context, agentKey string) (extra []actool.CoreTool, def DeferredInfo, cleanup func()) + +// The hook is set by server wiring while restored tasks may already be +// assembling tools on other goroutines, so every access goes through the lock. +var ( + toolAugmentMu sync.RWMutex + toolAugment ToolAugmentFunc +) + +// SetToolAugment installs (or, with nil, removes) the ToolAugment hook. +func SetToolAugment(f ToolAugmentFunc) { + toolAugmentMu.Lock() + toolAugment = f + toolAugmentMu.Unlock() +} + +// ToolAugment returns the current ToolAugment hook, or nil when none is set. +func ToolAugment() ToolAugmentFunc { + toolAugmentMu.RLock() + defer toolAugmentMu.RUnlock() + return toolAugment +} // AugmentTools returns base plus the agent's visible skill/MCP tools, the // DeferredInfo, and a cleanup func the caller must defer (closes MCP clients). @@ -42,10 +64,11 @@ func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool) cleanup = func() {} out = base ) - if ToolAugment != nil { + // Copy each hook once under its lock, then call it outside the lock. + if augment := ToolAugment(); augment != nil { var extra []actool.CoreTool var cl func() - extra, def, cl = ToolAugment(ctx, agentKey) + extra, def, cl = augment(ctx, agentKey) if cl != nil { cleanup = cl } @@ -57,8 +80,8 @@ func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool) // agent isn't bound to (or that are disabled) and swap in overridden // descriptions/schemas + default injection. MCP/skill/host tools have no row // and pass through untouched, so deferred/unlock wiring stays consistent. - if ToolResolve != nil { - out = ToolResolve(ctx, agentKey, out) + if resolve := ToolResolve(); resolve != nil { + out = resolve(ctx, agentKey, out) } // Tool RAG (off by default): lexical projection of the bound list; only removes. out = applyToolRAG(ctx, agentKey, out, def.Deferred, bm25Score) diff --git a/agent/assembly_test.go b/agent/assembly_test.go new file mode 100644 index 0000000..1675e54 --- /dev/null +++ b/agent/assembly_test.go @@ -0,0 +1,87 @@ +package agent + +import ( + "context" + "sync" + "testing" + + actool "github.com/Autumn-27/norma/tool" +) + +// Server wiring installs the hooks while restored tasks may already be assembling +// tools, rendering prompts and building settlement configs on other goroutines. +// Readers and writers run concurrently here, so under -race any unsynchronised +// hook access is reported regardless of scheduling. +func TestAgentHooksConcurrentSwap(t *testing.T) { + oldAugment, oldResolve, oldPolicy := ToolAugment(), ToolResolve(), FindingTrafficBindingEnabled() + oldPrompt, oldWrapup, oldWrapupTurns := PromptOverride(), WrapupOverride(), WrapupMaxTurnsOverride() + oldTimeout, oldTimeoutTurns := WrapupTaskTimeoutOverride(), WrapupTaskTimeoutTurnsOverride() + t.Cleanup(func() { + SetToolAugment(oldAugment) + SetToolResolve(oldResolve) + SetFindingTrafficBindingEnabled(oldPolicy) + SetPromptOverride(oldPrompt) + SetWrapupOverride(oldWrapup) + SetWrapupMaxTurnsOverride(oldWrapupTurns) + SetWrapupTaskTimeoutOverride(oldTimeout) + SetWrapupTaskTimeoutTurnsOverride(oldTimeoutTurns) + }) + augment := func(context.Context, string) ([]actool.CoreTool, DeferredInfo, func()) { + return nil, DeferredInfo{}, nil + } + resolve := func(_ context.Context, _ string, in []actool.CoreTool) []actool.CoreTool { return in } + policy := func() bool { return false } + text := func(string) (string, bool) { return "", false } + turns := func(string) (int, bool) { return 0, false } + setAll := func(on bool) { + if on { + SetToolAugment(augment) + SetToolResolve(resolve) + SetFindingTrafficBindingEnabled(policy) + SetPromptOverride(text) + SetWrapupOverride(text) + SetWrapupMaxTurnsOverride(turns) + SetWrapupTaskTimeoutOverride(text) + SetWrapupTaskTimeoutTurnsOverride(turns) + return + } + SetToolAugment(nil) + SetToolResolve(nil) + SetFindingTrafficBindingEnabled(nil) + SetPromptOverride(nil) + SetWrapupOverride(nil) + SetWrapupMaxTurnsOverride(nil) + SetWrapupTaskTimeoutOverride(nil) + SetWrapupTaskTimeoutTurnsOverride(nil) + } + + const rounds = 200 + var wg sync.WaitGroup + for range 4 { + wg.Go(func() { + for range rounds { + out, _, cleanup := AugmentTools(context.Background(), "worker", nil) + cleanup() + if len(out) != 0 { + t.Errorf("AugmentTools(nil base) = %d tools, want 0", len(out)) + return + } + if got := renderSystem("planner", "DEFAULT", PlannerVars{}); got != "DEFAULT" { + t.Errorf("renderSystem = %q, want the default", got) + return + } + if st := wrapupSettlementForTask("worker", nil, true); st.Prompt != WrapupDefault("worker") { + t.Errorf("settlement prompt = %q, want the default", st.Prompt) + return + } + } + }) + } + wg.Go(func() { + for range rounds { + setAll(true) + setAll(false) + } + }) + wg.Wait() +} diff --git a/agent/finding_recorder_test.go b/agent/finding_recorder_test.go index aca044f..13d274c 100644 --- a/agent/finding_recorder_test.go +++ b/agent/finding_recorder_test.go @@ -51,9 +51,9 @@ func TestReportFindingOptionalTrafficWithoutCapture(t *testing.T) { } func TestReportFindingAtomicContract(t *testing.T) { - old := FindingTrafficBindingEnabled - FindingTrafficBindingEnabled = func() bool { return true } - t.Cleanup(func() { FindingTrafficBindingEnabled = old }) + old := FindingTrafficBindingEnabled() + SetFindingTrafficBindingEnabled(func() bool { return true }) + t.Cleanup(func() { SetFindingTrafficBindingEnabled(old) }) d := testDB(t) defer d.Close() task, err := d.CreateTask("report contract", "fixture", nil, 0, 0) diff --git a/agent/finding_workflow.go b/agent/finding_workflow.go index 5231da4..1428646 100644 --- a/agent/finding_workflow.go +++ b/agent/finding_workflow.go @@ -3,6 +3,7 @@ package agent import ( "encoding/json" "fmt" + "sync" "github.com/Autumn-27/artex/db" actool "github.com/Autumn-27/norma/tool" @@ -13,10 +14,29 @@ const findingIDGuidance = "\n\n**Finding ID conventions**: finding_id is the sta // The server supplies the persisted setting. A missing setting/host is off. // Consulted at assembly and again on writes so an already-running session // cannot keep binding after the user switches the feature off. -var FindingTrafficBindingEnabled func() bool +// Guarded by a lock for the same reason as the tool hooks. +var ( + findingTrafficBindingMu sync.RWMutex + findingTrafficBindingHook func() bool +) + +// SetFindingTrafficBindingEnabled installs (or, with nil, removes) the hook. +func SetFindingTrafficBindingEnabled(f func() bool) { + findingTrafficBindingMu.Lock() + findingTrafficBindingHook = f + findingTrafficBindingMu.Unlock() +} + +// FindingTrafficBindingEnabled returns the current hook, or nil when none is set. +func FindingTrafficBindingEnabled() func() bool { + findingTrafficBindingMu.RLock() + defer findingTrafficBindingMu.RUnlock() + return findingTrafficBindingHook +} func findingTrafficBindingEnabled() bool { - return FindingTrafficBindingEnabled != nil && FindingTrafficBindingEnabled() + f := FindingTrafficBindingEnabled() + return f != nil && f() } // Applied after ToolResolve: user descriptions and prompts remain intact, while diff --git a/agent/finding_workflow_test.go b/agent/finding_workflow_test.go index fd2df65..c478f3d 100644 --- a/agent/finding_workflow_test.go +++ b/agent/finding_workflow_test.go @@ -10,19 +10,23 @@ import ( ) func TestFindingWorkflowSharedAssemblyAndSwitch(t *testing.T) { - oldPolicy, oldResolve, oldAugment := FindingTrafficBindingEnabled, ToolResolve, ToolAugment - t.Cleanup(func() { FindingTrafficBindingEnabled, ToolResolve, ToolAugment = oldPolicy, oldResolve, oldAugment }) - ToolAugment = nil + oldPolicy, oldResolve, oldAugment := FindingTrafficBindingEnabled(), ToolResolve(), ToolAugment() + t.Cleanup(func() { + SetFindingTrafficBindingEnabled(oldPolicy) + SetToolResolve(oldResolve) + SetToolAugment(oldAugment) + }) + SetToolAugment(nil) on := false - FindingTrafficBindingEnabled = func() bool { return on } + SetFindingTrafficBindingEnabled(func() bool { return on }) ts := NewToolSet(nil, "fixture") - ToolResolve = func(_ context.Context, _ string, base []actool.CoreTool) []actool.CoreTool { + SetToolResolve(func(_ context.Context, _ string, base []actool.CoreTool) []actool.CoreTool { out := make([]actool.CoreTool, len(base)) for i, tool := range base { out[i] = DecorateTool(tool, "CUSTOM DESCRIPTION", tool.InputSchema()) } return out - } + }) for _, role := range []string{"worker", "planner", "mainagent", "auto", "pentest", "custom-agent"} { for _, enabled := range []bool{false, true, false} { on = enabled @@ -59,7 +63,7 @@ func TestFindingWorkflowSharedAssemblyAndSwitch(t *testing.T) { } } on = true - ToolResolve = func(context.Context, string, []actool.CoreTool) []actool.CoreTool { return nil } + SetToolResolve(func(context.Context, string, []actool.CoreTool) []actool.CoreTool { return nil }) out, def, cleanup := AugmentTools(t.Context(), "planner", []actool.CoreTool{ts.addFinding()}) defer cleanup() if len(out) != 0 || def.FindingGuidance != "" { diff --git a/agent/prompt.go b/agent/prompt.go index f4aedca..3676ce6 100644 --- a/agent/prompt.go +++ b/agent/prompt.go @@ -2,6 +2,7 @@ package agent import ( "bytes" + "sync" "text/template" "time" ) @@ -10,7 +11,25 @@ import ( // key and whether one exists. The server wires it to the PG agent_prompts table. // When nil or no override exists, agents use their built-in default prompt — so // behavior is identical until a user edits a prompt in the UI. -var PromptOverride func(agentKey string) (string, bool) +// Guarded by a lock: server wiring can run while agents render prompts. +var ( + promptOverrideMu sync.RWMutex + promptOverride func(agentKey string) (string, bool) +) + +// SetPromptOverride installs (or, with nil, removes) the PromptOverride hook. +func SetPromptOverride(f func(agentKey string) (string, bool)) { + promptOverrideMu.Lock() + promptOverride = f + promptOverrideMu.Unlock() +} + +// PromptOverride returns the current hook, or nil when none is set. +func PromptOverride() func(agentKey string) (string, bool) { + promptOverrideMu.RLock() + defer promptOverrideMu.RUnlock() + return promptOverride +} // Prompt-variable structs — fields mirror each agent's catalog (docs §5a) so a // user template referencing a catalog variable renders; referencing anything else @@ -38,8 +57,8 @@ func nowStr() string { return time.Now().Format("2006-01-02 15:04:05 MST") } // the DB body. func renderSystem(agentKey, def string, vars any) string { tmpl := def - if PromptOverride != nil { - if t, ok := PromptOverride(agentKey); ok && t != "" { + if override := PromptOverride(); override != nil { + if t, ok := override(agentKey); ok && t != "" { tmpl = t } } diff --git a/agent/prompt_injection_seed_test.go b/agent/prompt_injection_seed_test.go index 51df7eb..48f6a98 100644 --- a/agent/prompt_injection_seed_test.go +++ b/agent/prompt_injection_seed_test.go @@ -18,9 +18,9 @@ func TestPromptSeedsDoNotCarryInjectionBlock(t *testing.T) { // tail: the #49 block for the roles that read target output, the untrusted-data // rule for the planner and the cold-digest compactor. func TestInjectionTailSurvivesPromptOverride(t *testing.T) { - prev := PromptOverride - t.Cleanup(func() { PromptOverride = prev }) - PromptOverride = func(string) (string, bool) { return "BARE BODY", true } + prev := PromptOverride() + t.Cleanup(func() { SetPromptOverride(prev) }) + SetPromptOverride(func(string) (string, bool) { return "BARE BODY", true }) content := map[string]string{ "worker": workerSystem("", "", "/data", "/run"), @@ -64,8 +64,8 @@ func TestInjectionTailSurvivesPromptOverride(t *testing.T) { // A body that already carries the block (a DB row seeded before the move) does // not get it twice; the code defaults carry it exactly once. func TestInjectionTailNotDuplicated(t *testing.T) { - prev := PromptOverride - t.Cleanup(func() { PromptOverride = prev }) + prev := PromptOverride() + t.Cleanup(func() { SetPromptOverride(prev) }) render := func() map[string]string { return map[string]string{ @@ -78,16 +78,16 @@ func TestInjectionTailNotDuplicated(t *testing.T) { } } - PromptOverride = func(string) (string, bool) { + SetPromptOverride(func(string) (string, bool) { return "LEGACY BODY\n\n" + untrustedContentRules + "\n\nMore rules.", true - } + }) for role, sys := range render() { if n := strings.Count(sys, untrustedContentRules); n != 1 { t.Errorf("legacy %s: untrusted-content block rendered %d times, want 1", role, n) } } - PromptOverride = nil + SetPromptOverride(nil) for role, sys := range render() { if n := strings.Count(sys, untrustedContentRules); n != 1 { t.Errorf("default %s: untrusted-content block rendered %d times, want 1", role, n) diff --git a/agent/prompt_now_test.go b/agent/prompt_now_test.go index 26984cc..707f228 100644 --- a/agent/prompt_now_test.go +++ b/agent/prompt_now_test.go @@ -10,15 +10,15 @@ import ( // (chat) agent prompt referencing it renders the live server time each turn, rather // than failing template execution and falling back to DefaultAssistantPrompt. func TestChatNowVarRenders(t *testing.T) { - prev := PromptOverride - defer func() { PromptOverride = prev }() + prev := PromptOverride() + defer func() { SetPromptOverride(prev) }() - PromptOverride = func(key string) (string, bool) { + SetPromptOverride(func(key string) (string, bool) { if key == "tec_benchmark" { return "Current time: {{.Now}}", true } return "", false - } + }) out := chatSystem("tec_benchmark", "/app/data", "/tmp/x") if !strings.Contains(out, "Current time: ") { @@ -34,12 +34,12 @@ func TestChatNowVarRenders(t *testing.T) { // a custom prompt referencing it renders the server data root (s.m.dir), rather // than failing template execution and falling back to DefaultAssistantPrompt. func TestChatDataDirVarRenders(t *testing.T) { - prev := PromptOverride - defer func() { PromptOverride = prev }() + prev := PromptOverride() + defer func() { SetPromptOverride(prev) }() - PromptOverride = func(key string) (string, bool) { + SetPromptOverride(func(key string) (string, bool) { return "Data root: {{.DataDir}}", true - } + }) out := chatSystem("tec_benchmark", "/app/data", "/tmp/x") if !strings.Contains(out, "Data root: /app/data") { @@ -50,12 +50,12 @@ func TestChatDataDirVarRenders(t *testing.T) { // TestChatUnknownVarFallsBack verifies an out-of-catalog {{.X}} still degrades // safely to the default assistant prompt (never a half-rendered prompt). func TestChatUnknownVarFallsBack(t *testing.T) { - prev := PromptOverride - defer func() { PromptOverride = prev }() + prev := PromptOverride() + defer func() { SetPromptOverride(prev) }() - PromptOverride = func(key string) (string, bool) { + SetPromptOverride(func(key string) (string, bool) { return "References a nonexistent variable: {{.Bogus}}", true - } + }) out := chatSystem("whatever", "/app/data", "/tmp/x") if strings.Contains(out, "References a nonexistent variable") { diff --git a/agent/prompt_test.go b/agent/prompt_test.go index 9157027..6ced28c 100644 --- a/agent/prompt_test.go +++ b/agent/prompt_test.go @@ -6,27 +6,27 @@ import ( ) func TestRenderSystemOverrideAndFallback(t *testing.T) { - t.Cleanup(func() { PromptOverride = nil }) + t.Cleanup(func() { SetPromptOverride(nil) }) // no override → built-in default - PromptOverride = nil + SetPromptOverride(nil) if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "g"}); got != "DEFAULT" { t.Fatalf("no override should give default, got %q", got) } // override → rendered with vars - PromptOverride = func(k string) (string, bool) { + SetPromptOverride(func(k string) (string, bool) { if k == "planner" { return "Goal:{{.Goal}} Scope:{{.Scope}}", true } return "", false - } + }) if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "capture X", Scope: "*.x.com"}); got != "Goal:capture X Scope:*.x.com" { t.Fatalf("override render: %q", got) } // override referencing a non-catalog var → execution error → fallback to default - PromptOverride = func(k string) (string, bool) { return "{{.NotInCatalog}}", true } + SetPromptOverride(func(k string) (string, bool) { return "{{.NotInCatalog}}", true }) if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "x"}); got != "DEFAULT" { t.Fatalf("bad var should fall back to default, got %q", got) } @@ -34,7 +34,7 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) { // full plannerSystem path: DB body [A] is honored, then the code-owned tail // [C] (intermediate artifact output rules) is ALWAYS appended — editing the body // can't drop it. - PromptOverride = func(k string) (string, bool) { return "PLANNER {{.Goal}}", true } + SetPromptOverride(func(k string) (string, bool) { return "PLANNER {{.Goal}}", true }) got := plannerSystem("capture X", "/data", "/data") if !strings.HasPrefix(got, "PLANNER capture X") { t.Fatalf("plannerSystem body not honored: %q", got) @@ -48,9 +48,9 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) { // the traffic_* tools exist), [C] artifact spec always present. The trafficTool // block is gated on the CA (arg 2), NOT on ProxyAddr — a global egress proxy // with capture off routes traffic but records nothing. - PromptOverride = func(k string) (string, bool) { + SetPromptOverride(func(k string) (string, bool) { return "{{if .ProxyAddr}}via proxy {{.ProxyAddr}}{{else}}manual{{end}}", true - } + }) recording := workerSystem("127.0.0.1:8080", "/ca.pem", "/data", "/data") if !strings.HasPrefix(recording, "via proxy 127.0.0.1:8080") { t.Fatalf("worker proxy branch body: %q", recording) diff --git a/agent/toolcatalog.go b/agent/toolcatalog.go index 276dbed..36fa00a 100644 --- a/agent/toolcatalog.go +++ b/agent/toolcatalog.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "sync" actool "github.com/Autumn-27/norma/tool" ) @@ -121,13 +122,34 @@ func BuiltinToolSeeds() []ToolSeed { return out } -// ToolResolve, if set, post-processes an agent's fully-assembled tool list against +// ToolResolveFunc is the ToolResolve hook. If set, it post-processes an agent's +// fully-assembled tool list against // the DB tools table: it drops tools not bound to this agent (or globally disabled) // and wraps the rest so the model sees the DB-overridden description/schema and // default inputs get injected. Tools with no matching DB row (MCP/skill/host tools // like traffic) pass through untouched. nil = tools unchanged. Wired in // server/assembly.go. -var ToolResolve func(ctx context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool +type ToolResolveFunc func(ctx context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool + +// Guarded like toolAugment: wiring can run while tasks assemble tools. +var ( + toolResolveMu sync.RWMutex + toolResolve ToolResolveFunc +) + +// SetToolResolve installs (or, with nil, removes) the ToolResolve hook. +func SetToolResolve(f ToolResolveFunc) { + toolResolveMu.Lock() + toolResolve = f + toolResolveMu.Unlock() +} + +// ToolResolve returns the current ToolResolve hook, or nil when none is set. +func ToolResolve() ToolResolveFunc { + toolResolveMu.RLock() + defer toolResolveMu.RUnlock() + return toolResolve +} // DecorateTool wraps t so Description()/InputSchema() report the DB overrides and // Call() injects scalar parameter defaults (from schema's "default" props) whenever diff --git a/agent/toolrag_test.go b/agent/toolrag_test.go index 51ceaad..4c73090 100644 --- a/agent/toolrag_test.go +++ b/agent/toolrag_test.go @@ -82,12 +82,17 @@ func TestToolRAGCoreKeptAndKRespected(t *testing.T) { // The projection runs after ToolResolve, so a tool the binding/enabled filter // dropped is never selected even when it is the best lexical match. func TestToolRAGNeverSelectsUnboundOrDisabled(t *testing.T) { - oldResolve, oldAugment, oldPolicy := ToolResolve, ToolAugment, FindingTrafficBindingEnabled - t.Cleanup(func() { ToolResolve, ToolAugment, FindingTrafficBindingEnabled = oldResolve, oldAugment, oldPolicy }) - ToolAugment, FindingTrafficBindingEnabled = nil, nil + oldResolve, oldAugment, oldPolicy := ToolResolve(), ToolAugment(), FindingTrafficBindingEnabled() + t.Cleanup(func() { + SetToolResolve(oldResolve) + SetToolAugment(oldAugment) + SetFindingTrafficBindingEnabled(oldPolicy) + }) + SetToolAugment(nil) + SetFindingTrafficBindingEnabled(nil) t.Setenv(toolRAGEnv, "on") t.Setenv(toolRAGTopKEnv, "2") - ToolResolve = func(_ context.Context, _ string, in []actool.CoreTool) []actool.CoreTool { + SetToolResolve(func(_ context.Context, _ string, in []actool.CoreTool) []actool.CoreTool { var out []actool.CoreTool for _, tool := range in { if tool.Name() != "sql_injection_exploit" { @@ -95,7 +100,7 @@ func TestToolRAGNeverSelectsUnboundOrDisabled(t *testing.T) { } } return out - } + }) base := append(ragCatalog(25), ragTool("sql_injection_exploit", "Exploit SQL injection in web login form", "url")) out, _, cleanup := AugmentTools(WithToolQuery(t.Context(), ragQuery), "worker", base) defer cleanup() diff --git a/agent/wrapup.go b/agent/wrapup.go index 8cb248f..5429c92 100644 --- a/agent/wrapup.go +++ b/agent/wrapup.go @@ -2,6 +2,7 @@ package agent import ( "strings" + "sync" "github.com/Autumn-27/norma/harness" ) @@ -17,15 +18,47 @@ import ( // [prompt body] is editable; which tools are disabled and how many turns the // wrap-up itself gets are code-fixed policy. -// WrapupOverride, if set, returns the stored wrap-up prompt for an agent key and +// wrapupHooksMu guards the four wrap-up hooks in this file: server wiring can run +// while agents build their settlement config on other goroutines. +var wrapupHooksMu sync.RWMutex + +// wrapupOverride, if set, returns the stored wrap-up prompt for an agent key and // whether a non-empty one exists. Wired by the server to the agents table (like // PromptOverride for system prompts). nil / empty → the built-in default is used. -var WrapupOverride func(agentKey string) (string, bool) +var wrapupOverride func(agentKey string) (string, bool) -// WrapupMaxTurnsOverride, if set, returns the admin-configured turn budget for the +// SetWrapupOverride installs (or, with nil, removes) the WrapupOverride hook. +func SetWrapupOverride(f func(agentKey string) (string, bool)) { + wrapupHooksMu.Lock() + wrapupOverride = f + wrapupHooksMu.Unlock() +} + +// WrapupOverride returns the current hook, or nil when none is set. +func WrapupOverride() func(agentKey string) (string, bool) { + wrapupHooksMu.RLock() + defer wrapupHooksMu.RUnlock() + return wrapupOverride +} + +// wrapupMaxTurnsOverride, if set, returns the admin-configured turn budget for the // wrap-up phase of an agent and whether a positive one exists. Wired to the agents // table. nil / ≤0 → the built-in per-agent default (wrapupTurnDefaults) is used. -var WrapupMaxTurnsOverride func(agentKey string) (int, bool) +var wrapupMaxTurnsOverride func(agentKey string) (int, bool) + +// SetWrapupMaxTurnsOverride installs (or, with nil, removes) the WrapupMaxTurnsOverride hook. +func SetWrapupMaxTurnsOverride(f func(agentKey string) (int, bool)) { + wrapupHooksMu.Lock() + wrapupMaxTurnsOverride = f + wrapupHooksMu.Unlock() +} + +// WrapupMaxTurnsOverride returns the current hook, or nil when none is set. +func WrapupMaxTurnsOverride() func(agentKey string) (int, bool) { + wrapupHooksMu.RLock() + defer wrapupHooksMu.RUnlock() + return wrapupMaxTurnsOverride +} // Built-in default wrap-up prompts, keyed by agent key. worker reuses the // historically hardcoded settleWrapUpPrompt (defined in worker.go); planner and @@ -75,8 +108,8 @@ func WrapupTurnsDefault(agentKey string) int { // resolveWrapup returns the effective wrap-up prompt: the DB override (if set and // non-empty) over the built-in default. func resolveWrapup(agentKey string) string { - if WrapupOverride != nil { - if t, ok := WrapupOverride(agentKey); ok && strings.TrimSpace(t) != "" { + if override := WrapupOverride(); override != nil { + if t, ok := override(agentKey); ok && strings.TrimSpace(t) != "" { return t } } @@ -86,8 +119,8 @@ func resolveWrapup(agentKey string) string { // resolveWrapupTurns returns the effective wrap-up turn budget: a positive DB // override over the built-in per-agent default. func resolveWrapupTurns(agentKey string) int { - if WrapupMaxTurnsOverride != nil { - if v, ok := WrapupMaxTurnsOverride(agentKey); ok && v > 0 { + if override := WrapupMaxTurnsOverride(); override != nil { + if v, ok := override(agentKey); ok && v > 0 { return v } } @@ -117,11 +150,40 @@ func wrapupSettlement(agentKey string, disabledTools []string) *harness.Settleme // WrapupTaskTimeoutOverride / …TurnsOverride: DB overrides for the task-timeout // wrap-up prompt and turns (wired to agents.task_timeout_wrapup_prompt / // _max_turns, worker/planner only). +// Guarded by wrapupHooksMu, like the per-run hooks above. var ( - WrapupTaskTimeoutOverride func(agentKey string) (string, bool) - WrapupTaskTimeoutTurnsOverride func(agentKey string) (int, bool) + wrapupTaskTimeoutOverride func(agentKey string) (string, bool) + wrapupTaskTimeoutTurnsOverride func(agentKey string) (int, bool) ) +// SetWrapupTaskTimeoutOverride installs (or, with nil, removes) the hook. +func SetWrapupTaskTimeoutOverride(f func(agentKey string) (string, bool)) { + wrapupHooksMu.Lock() + wrapupTaskTimeoutOverride = f + wrapupHooksMu.Unlock() +} + +// WrapupTaskTimeoutOverride returns the current hook, or nil when none is set. +func WrapupTaskTimeoutOverride() func(agentKey string) (string, bool) { + wrapupHooksMu.RLock() + defer wrapupHooksMu.RUnlock() + return wrapupTaskTimeoutOverride +} + +// SetWrapupTaskTimeoutTurnsOverride installs (or, with nil, removes) the hook. +func SetWrapupTaskTimeoutTurnsOverride(f func(agentKey string) (int, bool)) { + wrapupHooksMu.Lock() + wrapupTaskTimeoutTurnsOverride = f + wrapupHooksMu.Unlock() +} + +// WrapupTaskTimeoutTurnsOverride returns the current hook, or nil when none is set. +func WrapupTaskTimeoutTurnsOverride() func(agentKey string) (int, bool) { + wrapupHooksMu.RLock() + defer wrapupHooksMu.RUnlock() + return wrapupTaskTimeoutTurnsOverride +} + var taskTimeoutWrapupDefaults = map[string]string{ "worker": workerTaskTimeoutDefault, "planner": plannerTaskTimeoutDefault, @@ -141,8 +203,8 @@ func TaskTimeoutWrapupDefault(agentKey string) string { // string means this agent has no task-timeout prompt (not worker/planner) — the // caller should then fall back to the per-run prompt. func resolveTaskTimeoutWrapup(agentKey string) string { - if WrapupTaskTimeoutOverride != nil { - if t, ok := WrapupTaskTimeoutOverride(agentKey); ok && strings.TrimSpace(t) != "" { + if override := WrapupTaskTimeoutOverride(); override != nil { + if t, ok := override(agentKey); ok && strings.TrimSpace(t) != "" { return t } } @@ -150,8 +212,8 @@ func resolveTaskTimeoutWrapup(agentKey string) string { } func resolveTaskTimeoutTurns(agentKey string) int { - if WrapupTaskTimeoutTurnsOverride != nil { - if v, ok := WrapupTaskTimeoutTurnsOverride(agentKey); ok && v > 0 { + if override := WrapupTaskTimeoutTurnsOverride(); override != nil { + if v, ok := override(agentKey); ok && v > 0 { return v } } diff --git a/server/assembly.go b/server/assembly.go index 55fde2b..24bff8b 100644 --- a/server/assembly.go +++ b/server/assembly.go @@ -29,7 +29,7 @@ const mcpUntrustedNote = "\nOutput of this tool may contain attacker-controlled // 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()) { + agent.SetToolAugment(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 @@ -188,7 +188,7 @@ func wireAgentAugment(pg *db.DB, skillDir string, hostTools func() ([]actool.Cor UnlockSkill: unlockSkill, } return extra, def, cleanup - } + }) } func idSet(ids []int64) map[int64]bool { @@ -222,7 +222,7 @@ func seedPrompts(pg *db.DB) { // flag and, if kept, wrapped so the model sees the DB-overridden description/schema // and default params 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) } + agent.SetFindingTrafficBindingEnabled(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. @@ -248,7 +248,7 @@ func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) { // 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\nPrograms that need [interactive input] (msfconsole / interactive ssh login / REPLs like mysql, psql, python / password or yes/no prompts / nc reverse shells) must not use Bash (it has no stdin and would hang); open an interactive session with shell_open instead (and shell_close when done). One-shot, non-interactive commands still use Bash." - agent.ToolResolve = func(ctx context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool { + agent.SetToolResolve(func(ctx context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool { rows, err := pg.ListTools() if err != nil { log.Printf("[tools] failed to read tools table, passing through with code defaults: %v", err) @@ -333,7 +333,7 @@ func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) { } } return out - } + }) } func contains(ss []string, v string) bool { diff --git a/server/assembly_test.go b/server/assembly_test.go index 1eabeba..e25866c 100644 --- a/server/assembly_test.go +++ b/server/assembly_test.go @@ -49,7 +49,7 @@ func TestAssembleVisibleSkill(t *testing.T) { } wireAgentAugment(pg, skillDir, nil) - t.Cleanup(func() { agent.ToolAugment = nil }) + t.Cleanup(func() { agent.SetToolAugment(nil) }) ag, _ := pg.GetAgentByKey("planner") if ag == nil { @@ -65,7 +65,7 @@ func TestAssembleVisibleSkill(t *testing.T) { t.Cleanup(func() { _ = pg.SetAgentSkillVisibility(ag.ID, orig) }) // before: skill not visible → no Skill meta-tool - extra, _, cleanup := agent.ToolAugment(context.Background(), "planner") + extra, _, cleanup := agent.ToolAugment()(context.Background(), "planner") cleanup() if hasSkillTool(extra) { t.Fatalf("expected no Skill meta-tool before the skill is made visible") @@ -78,7 +78,7 @@ func TestAssembleVisibleSkill(t *testing.T) { t.Cleanup(func() { _ = pg.ToggleSkillVisibility(ag.ID, skillName, false) }) // after: the Skill meta-tool is packed in - extra, _, cleanup = agent.ToolAugment(context.Background(), "planner") + extra, _, cleanup = agent.ToolAugment()(context.Background(), "planner") defer cleanup() if !hasSkillTool(extra) { t.Fatalf("expected the Skill meta-tool after making the skill visible") diff --git a/server/fact_source_schema_test.go b/server/fact_source_schema_test.go index b180ec9..a677658 100644 --- a/server/fact_source_schema_test.go +++ b/server/fact_source_schema_test.go @@ -24,7 +24,7 @@ func TestRecordFactSourceSchemaUpgrade(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) - t.Cleanup(func() { agent.ToolResolve = nil }) + t.Cleanup(func() { agent.SetToolResolve(nil) }) pg := m.pg old, err := pg.GetTool("record_fact") @@ -96,7 +96,7 @@ func TestRecordFactSourceSchemaUpgrade(t *testing.T) { // serves it to the worker. func resolvedRecordFactProps(t *testing.T) map[string]any { t.Helper() - tools := agent.ToolResolve(context.Background(), "worker", agent.NewToolSet(nil, "").WorkerTools()) + tools := agent.ToolResolve()(context.Background(), "worker", agent.NewToolSet(nil, "").WorkerTools()) for _, tl := range tools { if tl.Name() == "record_fact" { props, _ := tl.InputSchema()["properties"].(map[string]any) diff --git a/server/finding_retests_test.go b/server/finding_retests_test.go index 860b860..78a8170 100644 --- a/server/finding_retests_test.go +++ b/server/finding_retests_test.go @@ -55,16 +55,18 @@ func newRetestServer(t *testing.T) (*Server, int64) { if err != nil { t.Fatal(err) } - oldAugment, oldResolve, oldPrompt := agent.ToolAugment, agent.ToolResolve, agent.PromptOverride - agent.ToolAugment = func(context.Context, string) ([]actool.CoreTool, agent.DeferredInfo, func()) { + oldAugment, oldResolve, oldPrompt := agent.ToolAugment(), agent.ToolResolve(), agent.PromptOverride() + agent.SetToolAugment(func(context.Context, string) ([]actool.CoreTool, agent.DeferredInfo, func()) { return s.findingRetestTools(), agent.DeferredInfo{}, func() {} - } - agent.ToolResolve = nil - agent.PromptOverride = func(string) (string, bool) { return agent.RetesterDefaultPrompt, true } + }) + agent.SetToolResolve(nil) + agent.SetPromptOverride(func(string) (string, bool) { return agent.RetesterDefaultPrompt, true }) t.Cleanup(func() { cancel() waitRetestIdle(t, s) - agent.ToolAugment, agent.ToolResolve, agent.PromptOverride = oldAugment, oldResolve, oldPrompt + agent.SetToolAugment(oldAugment) + agent.SetToolResolve(oldResolve) + agent.SetPromptOverride(oldPrompt) _, _ = pg.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid) _, _ = pg.DeleteFinding(fid) pg.Close() diff --git a/server/server.go b/server/server.go index 4f017b1..7d170c9 100644 --- a/server/server.go +++ b/server/server.go @@ -192,7 +192,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi // Wire DB-stored prompt templates into the agents (new design plan §3.3 / §5a). With no // override row, agents keep their built-in defaults — behavior is unchanged. if m.pg != nil { - agent.PromptOverride = func(key string) (string, bool) { + agent.SetPromptOverride(func(key string) (string, bool) { a, err := m.pg.GetAgentByKey(key) if err != nil || a == nil { return "", false @@ -202,38 +202,38 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi return "", false } return t, true - } + }) // Wire DB-stored wrap-up (settlement) prompts. Empty column → built-in default. - agent.WrapupOverride = func(key string) (string, bool) { + agent.SetWrapupOverride(func(key string) (string, bool) { a, err := m.pg.GetAgentByKey(key) if err != nil || a == nil || a.WrapupPrompt == "" { return "", false } return a.WrapupPrompt, true - } + }) // Wire DB-stored wrap-up turn budgets. 0 / missing → built-in per-agent default. - agent.WrapupMaxTurnsOverride = func(key string) (int, bool) { + agent.SetWrapupMaxTurnsOverride(func(key string) (int, bool) { a, err := m.pg.GetAgentByKey(key) if err != nil || a == nil || a.WrapupMaxTurns <= 0 { return 0, false } return a.WrapupMaxTurns, true - } + }) // Wire DB-stored TASK-TIMEOUT wrap-up prompt/turns (worker/planner). Empty/0 → default. - agent.WrapupTaskTimeoutOverride = func(key string) (string, bool) { + agent.SetWrapupTaskTimeoutOverride(func(key string) (string, bool) { a, err := m.pg.GetAgentByKey(key) if err != nil || a == nil || a.TaskTimeoutWrapupPrompt == "" { return "", false } return a.TaskTimeoutWrapupPrompt, true - } - agent.WrapupTaskTimeoutTurnsOverride = func(key string) (int, bool) { + }) + agent.SetWrapupTaskTimeoutTurnsOverride(func(key string) (int, bool) { a, err := m.pg.GetAgentByKey(key) if err != nil || a == nil || a.TaskTimeoutWrapupMaxTurns <= 0 { return 0, false } return a.TaskTimeoutWrapupMaxTurns, true - } + }) wireAgentAugment(m.pg, s.skillDir, s.hostTools) // wire visible skills/MCP + traffic/orchestration host tools into the agent toolset domainReg := buildDomainReg(m.Assets()) wireTools(m.pg, domainReg) // built-in tools table: filter per agent + override description/schema + inject defaults diff --git a/server/task_scope_description_test.go b/server/task_scope_description_test.go index d943385..2dee374 100644 --- a/server/task_scope_description_test.go +++ b/server/task_scope_description_test.go @@ -23,7 +23,7 @@ func TestTaskScopeDescriptionUpgrade(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) - t.Cleanup(func() { agent.ToolResolve = nil }) + t.Cleanup(func() { agent.SetToolResolve(nil) }) pg := m.pg var want string @@ -84,7 +84,7 @@ func TestTaskScopeDescriptionUpgrade(t *testing.T) { // ToolResolve serves it to the planner. func resolvedTaskScopeDescription(t *testing.T) string { t.Helper() - tools := agent.ToolResolve(context.Background(), "planner", agent.NewToolSet(nil, "").PlannerTools()) + tools := agent.ToolResolve()(context.Background(), "planner", agent.NewToolSet(nil, "").PlannerTools()) for _, tl := range tools { if tl.Name() == "add_task_scope" { return tl.Description() diff --git a/server/tools_wire_test.go b/server/tools_wire_test.go index 3e00bfb..4fa1323 100644 --- a/server/tools_wire_test.go +++ b/server/tools_wire_test.go @@ -33,7 +33,7 @@ func TestWireTools(t *testing.T) { defer pg.Close() wireTools(pg, nil) // nil domainReg: test only covers filter/decoration, not injection - t.Cleanup(func() { agent.ToolResolve = nil }) + t.Cleanup(func() { agent.SetToolResolve(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 @@ -52,7 +52,7 @@ func TestWireTools(t *testing.T) { base := append(ts.WorkerTools(), actool.DefaultTools()...) ctx := context.Background() - worker := names(agent.ToolResolve(ctx, "worker", base)) + worker := names(agent.ToolResolve()(ctx, "worker", base)) if _, ok := worker["record_fact"]; !ok { t.Error("worker lost record_fact") } @@ -65,7 +65,7 @@ func TestWireTools(t *testing.T) { 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)) + planner := names(agent.ToolResolve()(ctx, "planner", base)) if _, ok := planner["record_fact"]; ok { t.Error("planner should not get record_fact (not bound)") } @@ -85,7 +85,7 @@ func TestWireTools(t *testing.T) { t.Fatal(err) } - worker2 := names(agent.ToolResolve(ctx, "worker", base)) + worker2 := names(agent.ToolResolve()(ctx, "worker", base)) got := worker2["record_fact"] if got == nil { t.Fatal("record_fact missing after edit")