Merge pull request #227 from BaaRai/fix/agent-hook-race

agent: synchronise package-level hooks (data races under -race)

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)

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()
+}

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)

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

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 != "" {

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
 		}
 	}

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)

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") {

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)

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

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()

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
 		}
 	}

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 {

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")

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)

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()

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

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()

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")