From 37f2573a535b066f7dd4605d0d59d0c77d3f385a Mon Sep 17 00:00:00 2001 From: BaaRai <130491490+BaaRai@users.noreply.github.com> Date: Sat, 14 Nov 2026 13:46:03 +0100 Subject: [PATCH 1/2] agent: guard ToolAugment/ToolResolve hooks against concurrent access server wiring (wireAgentAugment, wireTools) assigned the package-level hooks agent.ToolAugment and agent.ToolResolve while restored tasks could already be running AugmentTools on other goroutines. go test -race ./server reported this as a DATA RACE between agent.AugmentTools and server.wireAgentAugment. Each hook is now an unexported variable behind a sync.RWMutex, with a setter (SetToolAugment, SetToolResolve) and a getter (ToolAugment(), ToolResolve()). AugmentTools copies the function value through the getter and calls it outside the lock. A nil hook still means no augmentation / no resolution, as before. FindingTrafficBindingEnabled is wired in the same server/assembly.go code and read on the same AugmentTools path, so it gets the same treatment. Adds TestAugmentToolsHooksConcurrentSwap, which runs AugmentTools while the setters swap the hooks, so -race catches an unsynchronised access without depending on test ordering. --- agent/assembly.go | 37 +++++++++++++++---- agent/assembly_test.go | 52 +++++++++++++++++++++++++++ agent/finding_recorder_test.go | 6 ++-- agent/finding_workflow.go | 24 +++++++++++-- agent/finding_workflow_test.go | 18 ++++++---- agent/toolcatalog.go | 26 ++++++++++++-- agent/toolrag_test.go | 15 +++++--- server/assembly.go | 10 +++--- server/assembly_test.go | 6 ++-- server/fact_source_schema_test.go | 4 +-- server/finding_retests_test.go | 12 ++++--- server/task_scope_description_test.go | 4 +-- server/tools_wire_test.go | 8 ++--- 13 files changed, 175 insertions(+), 47 deletions(-) create mode 100644 agent/assembly_test.go 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..bdaa6da --- /dev/null +++ b/agent/assembly_test.go @@ -0,0 +1,52 @@ +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 on other goroutines. Readers and writers run concurrently here, so under +// -race any unsynchronised hook access is reported regardless of scheduling. +func TestAugmentToolsHooksConcurrentSwap(t *testing.T) { + oldAugment, oldResolve, oldPolicy := ToolAugment(), ToolResolve(), FindingTrafficBindingEnabled() + t.Cleanup(func() { + SetToolAugment(oldAugment) + SetToolResolve(oldResolve) + SetFindingTrafficBindingEnabled(oldPolicy) + }) + 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 } + + 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 + } + } + }) + } + wg.Go(func() { + for range rounds { + SetToolAugment(augment) + SetToolResolve(resolve) + SetFindingTrafficBindingEnabled(policy) + SetToolAugment(nil) + SetToolResolve(nil) + SetFindingTrafficBindingEnabled(nil) + } + }) + 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/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/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..b5d6f39 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.SetToolResolve(nil) agent.PromptOverride = 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.PromptOverride = 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/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") From 54ff51177d96bfc15809fea12be33d939d204a0e Mon Sep 17 00:00:00 2001 From: BaaRai <130491490+BaaRai@users.noreply.github.com> Date: Sat, 14 Nov 2026 14:19:13 +0100 Subject: [PATCH 2/2] agent: guard prompt and wrap-up override hooks against concurrent access PromptOverride, WrapupOverride, WrapupMaxTurnsOverride, WrapupTaskTimeoutOverride and WrapupTaskTimeoutTurnsOverride are wired in server.New while restored tasks may already render prompts and build settlement configs on other goroutines. go test -race ./server reported races between server.New (server.go:195/207/215) and agent.renderSystem / resolveWrapup / resolveWrapupTurns. Same treatment as the tool hooks: unexported variables behind a sync.RWMutex (one for the prompt hook, one shared by the four wrap-up hooks), Set* setters and getters, and readers that copy the function value under the lock and call it outside. A nil hook still means the built-in default. TestAgentHooksConcurrentSwap now covers all eight hooks. --- agent/assembly_test.go | 53 ++++++++++++++--- agent/prompt.go | 25 +++++++- agent/prompt_injection_seed_test.go | 16 ++--- agent/prompt_now_test.go | 24 ++++---- agent/prompt_test.go | 16 ++--- agent/wrapup.go | 90 ++++++++++++++++++++++++----- server/finding_retests_test.go | 6 +- server/server.go | 20 +++---- 8 files changed, 183 insertions(+), 67 deletions(-) diff --git a/agent/assembly_test.go b/agent/assembly_test.go index bdaa6da..1675e54 100644 --- a/agent/assembly_test.go +++ b/agent/assembly_test.go @@ -9,20 +9,51 @@ import ( ) // Server wiring installs the hooks while restored tasks may already be assembling -// tools on other goroutines. Readers and writers run concurrently here, so under -// -race any unsynchronised hook access is reported regardless of scheduling. -func TestAugmentToolsHooksConcurrentSwap(t *testing.T) { +// 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 @@ -35,17 +66,21 @@ func TestAugmentToolsHooksConcurrentSwap(t *testing.T) { 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 { - SetToolAugment(augment) - SetToolResolve(resolve) - SetFindingTrafficBindingEnabled(policy) - SetToolAugment(nil) - SetToolResolve(nil) - SetFindingTrafficBindingEnabled(nil) + setAll(true) + setAll(false) } }) wg.Wait() 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/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/finding_retests_test.go b/server/finding_retests_test.go index b5d6f39..78a8170 100644 --- a/server/finding_retests_test.go +++ b/server/finding_retests_test.go @@ -55,18 +55,18 @@ func newRetestServer(t *testing.T) (*Server, int64) { if err != nil { t.Fatal(err) } - oldAugment, oldResolve, oldPrompt := agent.ToolAugment(), agent.ToolResolve(), agent.PromptOverride + 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.SetToolResolve(nil) - agent.PromptOverride = func(string) (string, bool) { return agent.RetesterDefaultPrompt, true } + agent.SetPromptOverride(func(string) (string, bool) { return agent.RetesterDefaultPrompt, true }) t.Cleanup(func() { cancel() waitRetestIdle(t, s) agent.SetToolAugment(oldAugment) agent.SetToolResolve(oldResolve) - agent.PromptOverride = oldPrompt + 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