MuyueWorkspace/internal/plugins/hooks.go

package plugins

import (
	"context"
	"encoding/json"
	"sync"

	"github.com/muyue/muyue/internal/agent"
)

type HookType string

const (
	BeforeToolCall      HookType = "before_tool_call"
	AfterToolCall       HookType = "after_tool_call"
	OnConversationStart HookType = "on_conversation_start"
	OnToolError         HookType = "on_tool_error"
)

type HookFunc func(ctx context.Context, payload HookPayload) error

type HookPayload struct {
	ToolName  string                 `json:"tool_name"`
	Arguments json.RawMessage        `json:"arguments,omitempty"`
	Response  *agent.ToolResponse    `json:"response,omitempty"`
	Error     string                 `json:"error,omitempty"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
}

type Hook struct {
	Type     HookType
	Plugin   string
	Priority int
	Fn       HookFunc
}

type HookRegistry struct {
	mu    sync.RWMutex
	hooks map[HookType][]Hook
}

func NewHookRegistry() *HookRegistry {
	return &HookRegistry{
		hooks: make(map[HookType][]Hook),
	}
}

func (hr *HookRegistry) Register(hookType HookType, pluginName string, priority int, fn HookFunc) {
	hr.mu.Lock()
	defer hr.mu.Unlock()

	h := Hook{
		Type:     hookType,
		Plugin:   pluginName,
		Priority: priority,
		Fn:       fn,
	}
	hr.hooks[hookType] = append(hr.hooks[hookType], h)

	for i := len(hr.hooks[hookType]) - 1; i > 0; i-- {
		if hr.hooks[hookType][i].Priority < hr.hooks[hookType][i-1].Priority {
			hr.hooks[hookType][i], hr.hooks[hookType][i-1] = hr.hooks[hookType][i-1], hr.hooks[hookType][i]
		}
	}
}

func (hr *HookRegistry) Fire(ctx context.Context, hookType HookType, payload HookPayload) error {
	hr.mu.RLock()
	hooks := make([]Hook, len(hr.hooks[hookType]))
	copy(hooks, hr.hooks[hookType])
	hr.mu.RUnlock()

	for _, h := range hooks {
		if err := h.Fn(ctx, payload); err != nil {
			return err
		}
	}
	return nil
}

func (hr *HookRegistry) RemoveByPlugin(pluginName string) {
	hr.mu.Lock()
	defer hr.mu.Unlock()

	for hookType := range hr.hooks {
		filtered := make([]Hook, 0, len(hr.hooks[hookType]))
		for _, h := range hr.hooks[hookType] {
			if h.Plugin != pluginName {
				filtered = append(filtered, h)
			}
		}
		hr.hooks[hookType] = filtered
	}
}