mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 04:30:55 +08:00
BaaRai
54ff51177d
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.
86 lines
3.3 KiB
Go
86 lines
3.3 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"sync"
|
|
"text/template"
|
|
"time"
|
|
)
|
|
|
|
// PromptOverride, if set, returns the stored system-prompt template for an agent
|
|
// 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.
|
|
// 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
|
|
// fails template execution and falls back to the built-in default.
|
|
type PlannerVars struct{ Goal, Scope, AssetSummary, DataDir, Now string }
|
|
type WorkerVars struct{ ProxyAddr, WorkerName, DataDir, Now string }
|
|
type MainVars struct{ Goal, AssetSummary, FindingsSummary, DataDir, Now string }
|
|
type GoalsVars struct{ EngagementDescription, DataDir, Now string }
|
|
|
|
// nowStr is the server-local wall-clock string exposed as the universal {{.Now}}
|
|
// prompt variable. renderSystem runs on every agent turn/round, so this is fresh
|
|
// each run — a prompt can subtract it from a fixed start stamp to reason about
|
|
// elapsed time (e.g. a timed benchmark's "last N hours" window).
|
|
func nowStr() string { return time.Now().Format("2006-01-02 15:04:05 MST") }
|
|
|
|
// renderSystem returns the rendered system-prompt BODY (section [A]) for agentKey.
|
|
// Precedence: the DB-stored template (if any) over the built-in default template
|
|
// (def). BOTH are Go templates now — the built-in default is seeded into the DB
|
|
// verbatim, so the two paths render identically until a user edits the prompt.
|
|
// Rendering always runs (def used to be pre-substituted plain text; it is now a
|
|
// {{.Var}} template like the DB one). On any render error we fall back to the
|
|
// default template, then to the raw default string — an agent never starts with a
|
|
// half-rendered prompt. Callers append the code-owned tail (trafficTool /
|
|
// intermediate artifact output rules) AFTER this, so those can't be edited away via
|
|
// the DB body.
|
|
func renderSystem(agentKey, def string, vars any) string {
|
|
tmpl := def
|
|
if override := PromptOverride(); override != nil {
|
|
if t, ok := override(agentKey); ok && t != "" {
|
|
tmpl = t
|
|
}
|
|
}
|
|
if out, err := renderTmpl(tmpl, vars); err == nil {
|
|
return out
|
|
}
|
|
// DB template broke (e.g. references an out-of-catalog var) → code default.
|
|
if out, err := renderTmpl(def, vars); err == nil {
|
|
return out
|
|
}
|
|
return def
|
|
}
|
|
|
|
func renderTmpl(tmpl string, vars any) (string, error) {
|
|
t, err := template.New("p").Option("missingkey=error").Parse(tmpl)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
var b bytes.Buffer
|
|
if err := t.Execute(&b, vars); err != nil {
|
|
return "", err
|
|
}
|
|
return b.String(), nil
|
|
}
|