fix(go): handle concurrent calls to Registry.ResolveAction()

go/internal/registry/registry.go

@@ -31,12 +31,12 @@ import (
 // Registry holds all registered actions and associated types,
 // and provides methods to register, query, and look up actions.
 type Registry struct {
-	mu      sync.Mutex
-	parent  api.Registry // parent registry for hierarchical lookups
-	actions map[string]api.Action
-	plugins map[string]api.Plugin
-	values  map[string]any // Values can truly be anything.
-
+	mu        sync.RWMutex
+	resolveMu sync.RWMutex
+	parent    api.Registry
+	actions   map[string]api.Action
+	plugins   map[string]api.Plugin
+	values    map[string]any // Values can truly be anything.
 	dotprompt *dotprompt.Dotprompt
 }
 
@@ -102,8 +102,8 @@ func (r *Registry) RegisterAction(key string, action api.Action) {
 // It first checks the current registry, then falls back to the parent if not found.
 // Returns nil if the plugin is not found in the registry hierarchy.
 func (r *Registry) LookupPlugin(name string) api.Plugin {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 
 	if plugin, ok := r.plugins[name]; ok {
 		return plugin
@@ -132,8 +132,8 @@ func (r *Registry) RegisterValue(name string, value any) {
 // It first checks the current registry, then falls back to the parent if not found.
 // Returns nil if the value is not found in the registry hierarchy.
 func (r *Registry) LookupValue(name string) any {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 
 	if value, ok := r.values[name]; ok {
 		return value
@@ -149,8 +149,8 @@ func (r *Registry) LookupValue(name string) any {
 // LookupAction returns the action for the given key.
 // It first checks the current registry, then falls back to the parent if not found.
 func (r *Registry) LookupAction(key string) api.Action {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 
 	if action, ok := r.actions[key]; ok {
 		return action
@@ -165,35 +165,47 @@ func (r *Registry) LookupAction(key string) api.Action {
 
 // ResolveAction looks up an action by key. If the action is not found, it attempts dynamic resolution.
 // Returns the action if found, or nil if not found.
+// This method is safe to call concurrently and uses a single mutex to serialize all resolution operations.
 func (r *Registry) ResolveAction(key string) api.Action {
 	action := r.LookupAction(key)
-	if action == nil {
-		typ, provider, name := api.ParseKey(key)
-		if typ == "" || name == "" {
-			slog.Debug("ResolveAction: failed to parse action key", "key", key)
-			return nil
-		}
-		plugins := r.ListPlugins()
-		for _, plugin := range plugins {
-			if dp, ok := plugin.(api.DynamicPlugin); ok && dp.Name() == provider {
-				resolvedAction := dp.ResolveAction(typ, name)
-				if resolvedAction != nil {
-					resolvedAction.Register(r)
-				}
-				break
+	if action != nil {
+		return action
+	}
+
+	r.resolveMu.Lock()
+	defer r.resolveMu.Unlock()
+
+	action = r.LookupAction(key)
+	if action != nil {
+		return action
+	}
+
+	typ, provider, name := api.ParseKey(key)
+	if typ == "" || name == "" {
+		slog.Debug("ResolveAction: failed to parse action key", "key", key)
+		return nil
+	}
+
+	plugins := r.ListPlugins()
+	for _, plugin := range plugins {
+		if dp, ok := plugin.(api.DynamicPlugin); ok && dp.Name() == provider {
+			resolvedAction := dp.ResolveAction(typ, name)
+			if resolvedAction != nil {
+				resolvedAction.Register(r)
 			}
+			break
 		}
-		action = r.LookupAction(key)
 	}
-	return action
+
+	return r.LookupAction(key)
 }
 
 // ListActions returns a list of all registered actions.
 // This includes actions from both the current registry and its parent hierarchy.
 // Child registry actions take precedence over parent actions with the same key.
 func (r *Registry) ListActions() []api.Action {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	var actions []api.Action
 	for _, v := range r.actions {
 		actions = append(actions, v)
@@ -203,8 +215,8 @@ func (r *Registry) ListActions() []api.Action {
 
 // ListPlugins returns a list of all registered plugins.
 func (r *Registry) ListPlugins() []api.Plugin {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	var plugins []api.Plugin
 	for _, p := range r.plugins {
 		plugins = append(plugins, p)
@@ -216,8 +228,8 @@ func (r *Registry) ListPlugins() []api.Plugin {
 // This includes values from both the current registry and its parent hierarchy.
 // Child registry values take precedence over parent values with the same key.
 func (r *Registry) ListValues() map[string]any {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 
 	allValues := make(map[string]any)
 
@@ -233,6 +245,8 @@ func (r *Registry) ListValues() map[string]any {
 
 // RegisterPartial adds the partial to the list of partials to the dotprompt instance
 func (r *Registry) RegisterPartial(name string, source string) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	if r.dotprompt == nil {
 		r.dotprompt = dotprompt.NewDotprompt(nil)
 	}
@@ -247,6 +261,8 @@ func (r *Registry) RegisterPartial(name string, source string) {
 
 // RegisterHelper adds a helper function to the dotprompt instance
 func (r *Registry) RegisterHelper(name string, fn any) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	if r.dotprompt == nil {
 		r.dotprompt = dotprompt.NewDotprompt(nil)
 	}
@@ -261,7 +277,7 @@ func (r *Registry) RegisterHelper(name string, fn any) {
 
 // Dotprompt returns the dotprompt instance.
 func (r *Registry) Dotprompt() *dotprompt.Dotprompt {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	return r.dotprompt
 }