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