Merge pull request #2531 from dgageot/board/simplifying-hooks-handling-7ba687cc

dgageot · web-flow · commit adcfafbed34d · 2026-04-27T15:43:41.000+02:00
refactor(hooks): simplify caching, dispatch flow, and notification helpers
diff --git a/pkg/hooks/executor.go b/pkg/hooks/executor.go
@@ -29,16 +29,12 @@ type Executor struct {
 // matcher is the compiled form of a [MatcherConfig]: an optional regex
 // pattern (nil means "match all") and the hooks to fire when it matches.
 type matcher struct {
-	raw     string
 	pattern *regexp.Regexp
 	hooks   []Hook
 }
 
 func (m *matcher) matches(toolName string) bool {
-	if m.raw == "" || m.raw == "*" {
-		return true
-	}
-	return m.pattern != nil && m.pattern.MatchString(toolName)
+	return m.pattern == nil || m.pattern.MatchString(toolName)
 }
 
 // hookResult is the outcome of a single hook invocation.
@@ -103,7 +99,7 @@ func compileMatchers(configs []MatcherConfig) []matcher {
 	}
 	out := make([]matcher, 0, len(configs))
 	for _, mc := range configs {
-		m := matcher{raw: mc.Matcher, hooks: mc.Hooks}
+		m := matcher{hooks: mc.Hooks}
 		if mc.Matcher != "" && mc.Matcher != "*" {
 			p, err := regexp.Compile("^(?:" + mc.Matcher + ")$")
 			if err != nil {
@@ -127,36 +123,17 @@ func (e *Executor) Has(event EventType) bool {
 // don't have to remember. Defaults [Input.Cwd] to the executor's
 // working directory when the caller didn't supply one.
 func (e *Executor) Dispatch(ctx context.Context, event EventType, input *Input) (*Result, error) {
-	matchers := e.events[event]
-	if len(matchers) == 0 {
+	hooks := e.hooksFor(event, input.ToolName)
+	if len(hooks) == 0 {
 		return &Result{Allowed: true}, nil
 	}
+
 	input.HookEventName = event
 	if input.Cwd == "" {
 		input.Cwd = e.workingDir
 	}
 
-	// Collect, filter by tool name, and dedup by (type, command, args).
-	// Dedup catches the common case of an explicit YAML hook overlapping
-	// a runtime auto-injected one (e.g. WithAddDate plus a user-authored
-	// add_date entry).
-	seen := make(map[string]bool)
-	var hooks []Hook
-	for _, m := range matchers {
-		if !m.matches(input.ToolName) {
-			continue
-		}
-		for _, h := range m.hooks {
-			key := dedupKey(h)
-			if !seen[key] {
-				seen[key] = true
-				hooks = append(hooks, h)
-			}
-		}
-	}
-	if len(hooks) == 0 {
-		return &Result{Allowed: true}, nil
-	}
+	slog.Debug("Executing hooks", "event", event, "session_id", input.SessionID, "count", len(hooks))
 
 	inputJSON, err := input.ToJSON()
 	if err != nil {
@@ -173,6 +150,29 @@ func (e *Executor) Dispatch(ctx context.Context, event EventType, input *Input)
 	return aggregate(results, event), nil
 }
 
+// hooksFor returns the deduplicated list of hooks that should run for
+// (event, toolName). Dedup by (type, command, args) catches the common
+// case of an explicit YAML hook overlapping a runtime auto-injected
+// one (e.g. WithAddDate plus a user-authored add_date entry).
+func (e *Executor) hooksFor(event EventType, toolName string) []Hook {
+	seen := make(map[string]bool)
+	var hooks []Hook
+	for _, m := range e.events[event] {
+		if !m.matches(toolName) {
+			continue
+		}
+		for _, h := range m.hooks {
+			key := dedupKey(h)
+			if seen[key] {
+				continue
+			}
+			seen[key] = true
+			hooks = append(hooks, h)
+		}
+	}
+	return hooks
+}
+
 // dedupKey returns a deterministic key identifying a hook by (type, command, args).
 func dedupKey(h Hook) string {
 	var b strings.Builder
diff --git a/pkg/runtime/hooks.go b/pkg/runtime/hooks.go
@@ -11,60 +11,50 @@ import (
 	"github.com/docker/docker-agent/pkg/session"
 )
 
-// hooksExec returns the cached [hooks.Executor] for a, building one on
-// first lookup. Returns nil when the agent has no user-configured hooks
-// and no agent-flag (AddDate / AddEnvironmentInfo / AddPromptFiles) maps
-// to a builtin. Callers can then short-circuit without paying for a
-// no-op dispatch.
+// buildHooksExecutors builds a [hooks.Executor] for every agent in the
+// team that has user-configured hooks or an agent-flag that maps to a
+// builtin (AddDate / AddEnvironmentInfo / AddPromptFiles). Agents with
+// no hooks have no entry; lookups fall through to nil so callers can
+// short-circuit cheaply.
 //
-// The cache is keyed by agent name. Entries (including the nil sentinel)
-// are stable for the lifetime of the runtime, so repeated dispatches
-// during a turn don't re-translate agent flags into builtin hook entries
-// or rebuild matcher tables.
+// Called once from [NewLocalRuntime] after r.workingDir, r.env and
+// r.hooksRegistry are finalized; the resulting map is read-only for
+// the lifetime of the runtime, so per-dispatch lookups don't need to
+// lock.
+func (r *LocalRuntime) buildHooksExecutors() {
+	r.hooksExecByAgent = make(map[string]*hooks.Executor)
+	for _, name := range r.team.AgentNames() {
+		a, err := r.team.Agent(name)
+		if err != nil {
+			continue
+		}
+		cfg := builtins.ApplyAgentDefaults(hooks.FromConfig(a.Hooks()), builtins.AgentDefaults{
+			AddDate:            a.AddDate(),
+			AddEnvironmentInfo: a.AddEnvironmentInfo(),
+			AddPromptFiles:     a.AddPromptFiles(),
+		})
+		if cfg == nil {
+			continue
+		}
+		r.hooksExecByAgent[name] = hooks.NewExecutorWithRegistry(cfg, r.workingDir, r.env, r.hooksRegistry)
+	}
+}
+
+// hooksExec returns the pre-built [hooks.Executor] for a, or nil when
+// the agent has no hooks (see [buildHooksExecutors]).
 func (r *LocalRuntime) hooksExec(a *agent.Agent) *hooks.Executor {
 	if a == nil {
 		return nil
 	}
-	name := a.Name()
-
-	r.hooksExecMu.RLock()
-	if exec, ok := r.hooksExecByAgent[name]; ok {
-		r.hooksExecMu.RUnlock()
-		return exec
-	}
-	r.hooksExecMu.RUnlock()
-
-	r.hooksExecMu.Lock()
-	defer r.hooksExecMu.Unlock()
-	// Re-check under the write lock to avoid double-build under contention.
-	if exec, ok := r.hooksExecByAgent[name]; ok {
-		return exec
-	}
-
-	cfg := builtins.ApplyAgentDefaults(hooks.FromConfig(a.Hooks()), builtins.AgentDefaults{
-		AddDate:            a.AddDate(),
-		AddEnvironmentInfo: a.AddEnvironmentInfo(),
-		AddPromptFiles:     a.AddPromptFiles(),
-	})
-
-	var exec *hooks.Executor
-	if cfg != nil {
-		exec = hooks.NewExecutorWithRegistry(cfg, r.workingDir, r.env, r.hooksRegistry)
-	}
-	if r.hooksExecByAgent == nil {
-		r.hooksExecByAgent = make(map[string]*hooks.Executor)
-	}
-	r.hooksExecByAgent[name] = exec
-	return exec
+	return r.hooksExecByAgent[a.Name()]
 }
 
 // dispatchHook is the common dispatch path shared by every hook
-// callsite: resolve the cached executor, short-circuit if no hook is
-// configured for event, then dispatch and emit any [Result.SystemMessage]
-// as a Warning event. Errors are logged at warn level and surfaced as
-// nil results so callers can use a single nil check to mean "nothing
-// useful came back" — covering the not-configured, no-agent, and
-// dispatch-failed cases uniformly.
+// callsite: resolve the pre-built executor, dispatch, and emit any
+// [Result.SystemMessage] as a Warning event. Errors are logged at warn
+// level and surfaced as nil results so callers can use a single nil
+// check to mean "nothing useful came back" — covering the
+// not-configured, no-agent, and dispatch-failed cases uniformly.
 //
 // events may be nil for fire-and-forget callsites (notification,
 // on_error, on_max_iterations, ...) where there's no Warning channel
@@ -79,11 +69,10 @@ func (r *LocalRuntime) dispatchHook(
 	events chan Event,
 ) *hooks.Result {
 	exec := r.hooksExec(a)
-	if exec == nil || !exec.Has(event) {
+	if exec == nil {
 		return nil
 	}
 
-	slog.Debug("Executing hooks", "event", event, "agent", a.Name(), "session_id", input.SessionID)
 	result, err := exec.Dispatch(ctx, event, input)
 	if err != nil {
 		slog.Warn("Hook execution failed", "event", event, "agent", a.Name(), "error", err)
@@ -150,41 +139,31 @@ func (r *LocalRuntime) executeStopHooks(ctx context.Context, sess *session.Sessi
 	}, events)
 }
 
-// executeNotificationHooks runs notification hooks when the agent emits
-// a user-facing notification. Hook output is informational — it does
-// not suppress or rewrite the notification.
-func (r *LocalRuntime) executeNotificationHooks(ctx context.Context, a *agent.Agent, sessionID, level, message string) {
-	if level != "error" && level != "warning" {
-		slog.Error("Invalid notification level", "level", level, "expected", "error|warning")
-		return
-	}
-	r.dispatchHook(ctx, a, hooks.EventNotification, &hooks.Input{
-		SessionID:           sessionID,
-		NotificationLevel:   level,
-		NotificationMessage: message,
-	}, nil)
+// notifyError fires both notification(level=error) and on_error in one
+// call. They're always emitted together (an error is always also a
+// user-facing notification), so collapsing them into one call expresses
+// intent more directly than firing two events at every callsite.
+func (r *LocalRuntime) notifyError(ctx context.Context, a *agent.Agent, sessionID, message string) {
+	r.notify(ctx, a, hooks.EventNotification, sessionID, "error", message)
+	r.notify(ctx, a, hooks.EventOnError, sessionID, "error", message)
 }
 
-// executeOnErrorHooks fires on_error when the runtime hits an error
-// during a turn (model failures, tool-call loops). Fires alongside the
-// broader notification event; on_error is the structured entry point
-// for users who want to react only to errors.
-func (r *LocalRuntime) executeOnErrorHooks(ctx context.Context, a *agent.Agent, sessionID, message string) {
-	r.dispatchHook(ctx, a, hooks.EventOnError, &hooks.Input{
-		SessionID:           sessionID,
-		NotificationLevel:   "error",
-		NotificationMessage: message,
-	}, nil)
+// notifyMaxIterations fires both notification(level=warning) and
+// on_max_iterations. Same rationale as [notifyError]: the two are
+// always emitted together when the iteration limit is reached.
+func (r *LocalRuntime) notifyMaxIterations(ctx context.Context, a *agent.Agent, sessionID, message string) {
+	r.notify(ctx, a, hooks.EventNotification, sessionID, "warning", message)
+	r.notify(ctx, a, hooks.EventOnMaxIterations, sessionID, "warning", message)
 }
 
-// executeOnMaxIterationsHooks fires on_max_iterations when the runtime
-// reaches its configured max_iterations limit. Fires alongside the
-// broader notification event; on_max_iterations is the structured entry
-// point for users who want to react only to that condition.
-func (r *LocalRuntime) executeOnMaxIterationsHooks(ctx context.Context, a *agent.Agent, sessionID, message string) {
-	r.dispatchHook(ctx, a, hooks.EventOnMaxIterations, &hooks.Input{
+// notify is the shared dispatch path for the (level, message)-shaped
+// hook events: notification, on_error, on_max_iterations. They all
+// take the same Input fields and are observational (no Result is
+// honored), so a single helper covers them all.
+func (r *LocalRuntime) notify(ctx context.Context, a *agent.Agent, event hooks.EventType, sessionID, level, message string) {
+	r.dispatchHook(ctx, a, event, &hooks.Input{
 		SessionID:           sessionID,
-		NotificationLevel:   "warning",
+		NotificationLevel:   level,
 		NotificationMessage: message,
 	}, nil)
 }
@@ -215,11 +194,10 @@ func (r *LocalRuntime) executeAfterLLMCallHooks(ctx context.Context, sess *sessi
 
 // executeOnUserInputHooks fires on_user_input when the runtime is about
 // to wait for the user (tool confirmation, elicitation, max iterations,
-// stream stopped). Resolves the agent from r.team itself so callsites
-// in code paths without an agent handle (like the elicitation handler)
-// stay short.
+// stream stopped). Resolves the agent itself so callsites in code paths
+// without an agent handle (like the elicitation handler) stay short.
 func (r *LocalRuntime) executeOnUserInputHooks(ctx context.Context, sessionID, logContext string) {
-	a, _ := r.team.Agent(r.CurrentAgentName())
+	a := r.CurrentAgent()
 	if a == nil {
 		return
 	}
diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go
@@ -252,8 +252,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				events <- MaxIterationsReached(runtimeMaxIterations)
 
 				maxIterMsg := fmt.Sprintf("Maximum iterations reached (%d)", runtimeMaxIterations)
-				r.executeNotificationHooks(ctx, a, sess.ID, "warning", maxIterMsg)
-				r.executeOnMaxIterationsHooks(ctx, a, sess.ID, maxIterMsg)
+				r.notifyMaxIterations(ctx, a, sess.ID, maxIterMsg)
 				r.executeOnUserInputHooks(ctx, sess.ID, "max iterations reached")
 
 				// In non-interactive mode (e.g. MCP server), auto-stop instead of
@@ -434,8 +433,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				telemetry.RecordError(ctx, err.Error())
 				errMsg := modelerrors.FormatError(err)
 				events <- Error(errMsg)
-				r.executeNotificationHooks(ctx, a, sess.ID, "error", errMsg)
-				r.executeOnErrorHooks(ctx, a, sess.ID, errMsg)
+				r.notifyError(ctx, a, sess.ID, errMsg)
 				streamSpan.End()
 				return
 			}
@@ -501,8 +499,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 						"This indicates a degenerate loop where the model is not making progress.",
 					loopDetector.consecutive, toolName)
 				events <- Error(errMsg)
-				r.executeNotificationHooks(ctx, a, sess.ID, "error", errMsg)
-				r.executeOnErrorHooks(ctx, a, sess.ID, errMsg)
+				r.notifyError(ctx, a, sess.ID, errMsg)
 				loopDetector.reset()
 				return
 			}
diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
@@ -206,14 +206,12 @@ type LocalRuntime struct {
 	// touching any process-wide state.
 	hooksRegistry *hooks.Registry
 
-	// hooksExecByAgent caches the per-agent [hooks.Executor], keyed by
-	// agent name. Building one requires translating agent flags into
-	// implicit builtin hook entries and compiling matchers — work we
-	// don't want to repeat on every dispatch within a turn. Entries are
-	// stable for the runtime's lifetime; a nil value caches the
-	// "no hooks configured" verdict so repeat lookups stay cheap.
+	// hooksExecByAgent holds the per-agent [hooks.Executor], keyed by
+	// agent name. Built once in [NewLocalRuntime.buildHooksExecutors]
+	// after team and runtime config are finalized; agents with no hooks
+	// have no entry, so [hooksExec] returns nil for them. Read-only after
+	// construction, so no locking is needed.
 	hooksExecByAgent map[string]*hooks.Executor
-	hooksExecMu      sync.RWMutex
 
 	// retryOnRateLimit enables retry-with-backoff for HTTP 429 (rate limit) errors
 	// when no fallback models are configured. When false (default), 429 errors are
@@ -394,6 +392,10 @@ func NewLocalRuntime(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
 	// RunStream on the same runtime (e.g. background agent sessions).
 	r.registerDefaultTools()
 
+	// Pre-build per-agent hook executors now that workingDir, env and
+	// the team are finalized. Read-only afterwards.
+	r.buildHooksExecutors()
+
 	slog.Debug("Creating new runtime", "agent", r.currentAgent, "available_agents", agents.Size())
 
 	return r, nil
