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.
This commit is contained in:
BaaRai
2026-11-14 14:19:13 +01:00
parent 37f2573a53
commit 54ff51177d
8 changed files with 183 additions and 67 deletions
+44 -9
View File
@@ -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()
+22 -3
View File
@@ -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
}
}
+8 -8
View File
@@ -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)
+12 -12
View File
@@ -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") {
+8 -8
View File
@@ -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)
+76 -14
View File
@@ -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
}
}
+3 -3
View File
@@ -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()
+10 -10
View File
@@ -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