HomeAgent/third_party/homeagent-sdk/sdk/plugin.go

package sdk

import (
	"sync"

	"gitcode.com/JianFeeeee/homeagent-sdk/meta"
)

// SDKVersion 是对外暴露的 SDK 版本号。
var SDKVersion = meta.Version

// Plugin is the interface every plugin must implement.
type Plugin interface {
	Name() string
	Start(sdk *PluginSDK) error
	Stop() error
}

// ToolHandler is a function that handles a tool call.
type ToolHandler func(args map[string]interface{}) (interface{}, error)

// StageHandler is a function that handles a pipeline stage event.
type StageHandler func(ctx *StageContext) error

// Stage represents a point in the message processing pipeline.
type Stage string

const (
	StageOnInput        Stage = "on_input"
	StagePreAction      Stage = "pre_action"
	StagePostAction     Stage = "post_action"
	StageBeforeToolcall Stage = "before_toolcall"
	StageAfterToolcall  Stage = "after_toolcall"
	StageBeforeOutput   Stage = "before_output"
	StageAfterOutput    Stage = "after_output"
)

// ChannelDef 描述通道在记忆计算层的行为，与 ToolDef.NoMemory/Cleaner 语义一致。
// NoMemory: 此通道输入/输出不参与记忆计算（向量化/关键词提取/蒸馏），但原文保留在上下文中
// Cleaner:  计算层过滤函数，不改原文；仅在向量化/jieba/蒸馏/存档提取关键词时调用
type ChannelDef struct {
	NoMemory bool
	Cleaner  func(string) string
}

// StageContext provides context for stage handlers.
type StageContext struct {
	mu               sync.RWMutex
	RawMessage       string
	UserID           string
	GroupID          string
	ContextMsgs      []map[string]interface{}
	LLMText          string
	ReasoningContent string
	TokenUsage       map[string]int
	ToolCalls        []ToolCall
	ToolResults      []ToolResult
	FinalText        string
	Response         *string
	Phase            Stage
	Memory           []MemItem
	NoMemory         bool
	Extra            map[string]interface{}
	Errors           []string // 阶段处理过程中的错误信息
}

func (c *StageContext) RLock()   { c.mu.RLock() }
func (c *StageContext) RUnlock() { c.mu.RUnlock() }
func (c *StageContext) Lock()    { c.mu.Lock() }
func (c *StageContext) Unlock()  { c.mu.Unlock() }
func (c *StageContext) IsResponded() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.Response != nil
}

// MemItem represents a memory item in stage context.
type MemItem struct {
	Role    string  `json:"role"`
	Content string  `json:"content"`
	Score   float64 `json:"score"`
}

// ToolCall represents a model's request to call a tool.
type ToolCall struct {
	ID        string                 `json:"id"`
	Name      string                 `json:"name"`
	Plugin    string                 `json:"plugin,omitempty"`
	Arguments map[string]interface{} `json:"arguments"`
}

// ToolResult represents the result of a tool call.
type ToolResult struct {
	CallID  string      `json:"call_id"`
	Name    string      `json:"name"`
	Plugin  string      `json:"plugin,omitempty"`
	Success bool        `json:"success"`
	Result  interface{} `json:"result"`
}

// ToolDef describes a tool that the plugin exposes.
type ToolDef struct {
	Name        string                 `json:"name"`
	Plugin      string                 `json:"plugin,omitempty"`
	Description string                 `json:"description"`
	Parameters  map[string]interface{} `json:"parameters"`
	NoMemory    bool                   `json:"no_memory,omitempty"` // 此工具输出不参与记忆计算，但原文保留
	Cleaner     func(string) string    `json:"-"`                   // 计算层过滤函数，不改原文；仅在向量化/jieba/蒸馏时调用
}

// IOInjector provides methods for injecting input and interrupts into the agent pipeline.
// All methods accept (source, channel) where channel is the target output channel
// for routing the agent's response.
type IOInjector interface {
	InjectInterruptText(source, channel, text string)
	InjectText(source, channel, text string)
	InjectTextNoMemory(source, channel, text string)
	// InjectInputSync 注入输入事件并同步等待 agent 回复，返回回复文本（无回复时返回空串）。
	// 用于通道消息的完整闭环：收到入站 → agent 处理 → 回复取回 → 送回通道。
	InjectInputSync(source, channel, text string) string
	// SetToolBlocks 插件工具注入多模态内容块（image_url/audio_url），内核在下一条
	// tool message 的 content 数组里带上这些块，让模型在后续轮次看到图/听到音频。
	SetToolBlocks(blocks []ContentBlock)
	InjectInputMedia(source, channel, text string, blocks []ContentBlock)
	InjectInputMediaSync(source, channel, text string, blocks []ContentBlock) string
	InjectInterruptMedia(source, channel, text string, blocks []ContentBlock)
}

// EventType identifies the kind of system event.
type EventType string

const (
	EventRawInput      EventType = "raw_input"
	EventAgentOutput   EventType = "agent_output"
	EventAgentLLMChain EventType = "agent_llm_chain"
	EventToolCall      EventType = "tool_call"
	EventReasoning     EventType = "reasoning"
	EventStage         EventType = "stage"
	EventSystem        EventType = "system"

	// 流式增量事件（token 级）：核心 process() 流式化后每收到一个增量块发布。
	// 客户端可选订做真逐 token 渲染；聚合事件仍照常发布，旧订阅者不受影响。
	EventReasoningDelta EventType = "reasoning_delta"
	EventContentDelta   EventType = "content_delta"
)

// Event represents a system event published by the kernel.
type Event struct {
	Type      EventType              `json:"type"`
	Source    string                 `json:"source"`
	Payload   map[string]interface{} `json:"payload"`
	Timestamp int64                  `json:"timestamp"`
}

// EventHandler processes a system event.
type EventHandler func(evt *Event)

// EventSubscriber allows plugins to subscribe to kernel events.
// This is a restricted interface: plugins can subscribe but the kernel
// controls which events are delivered.
type EventSubscriber interface {
	Subscribe(eventType EventType, handler EventHandler) func()
}

// PluginMgrAPI 提供插件管理能力（外部插件可调用）。
// 由 bridge 注入 dispatch 实现，走 C ABI CORE_PLUGIN_RELOAD_ONE 等。
type PluginMgrAPI interface {
	// ReloadOne 重载单个插件（停止后重新加载）。
	ReloadOne(name string) error
	// ListLoadedPlugins 列出已加载插件。
	ListLoadedPlugins() []string
	// IsPluginDisabled 查询插件是否被禁用。
	IsPluginDisabled(name string) bool
}

// StageScope controls which events a stage handler receives.
type StageScope int

const (
	// StageScopeGlobal receives all stage events (default).
	StageScopeGlobal StageScope = 0
	// StageScopeOwnTools only receives events for this plugin's own tool calls
	// (before_toolcall / after_toolcall only). Other stages degrade to global.
	StageScopeOwnTools StageScope = 1
)

// ToolRegistrar registers a tool dynamically.
type ToolRegistrar func(name string, def ToolDef, handler ToolHandler) error

// StageRegistrar registers a stage handler.
type StageRegistrar func(stage Stage, handler StageHandler)

// APIRegistrar registers a plugin API for external access.
type APIRegistrar func(name string) error

// InputChannelRegistrar registers an input channel with its memory behavior.
type InputChannelRegistrar func(name string, def ChannelDef) error

// OutputChannelRegistrar registers an output channel that the output_send tool can use.
type OutputChannelRegistrar func(name string, caps int, desc string, def ChannelDef, handler ToolHandler) error

// Output capability flags
const (
	CapText       = 1
	CapFile       = 2
	CapImage      = 4
	CapAudio      = 8
	CapStructured = 16
)

// PluginSDK is the main API surface provided to plugins at runtime.
// It wraps tool registration, settings, memory, knowledge, LLM, and IO injection.
type PluginSDK struct {
	name      string
	regTool   ToolRegistrar
	regStage  StageRegistrar
	regAPI    APIRegistrar
	regOutput OutputChannelRegistrar
	regInput  InputChannelRegistrar
	io        IOInjector
	mem       MemoryAPI
	textMem   TextMemoryAPI
	docMem    DocMemoryAPI
	know      KnowledgeAPI
	llm       LLMAPI
	sett      SettingsAPI
	social    SocialAPI
	events    EventSubscriber
	plgMgr    PluginMgrAPI

	// apiMu 保护上面这些由内核注入的 API 字段，以及 autoRestart。
	//
	// 这些字段的写方与读方天然跨 goroutine：
	//   - 写方是内核（加载/重载插件时注入 API）与插件自己（SetAutoRestart）；
	//   - 读方是插件在 Start() 里起的后台 goroutine（轮询、监听、定时器
	//     都要拿 injector 往管道里注消息），以及内核 registry —— 它在
	//     另一个 goroutine 读 AutoRestart() 决定崩溃后是否重启。
	// SetAutoRestart 的文档用法本身就是「连接建立后再决定能否自动重启」，
	// 而连接建立通常发生在后台 goroutine 里，于是这对读写必然并发。
	//
	// sdk/stress_test.go 的 -race 实测确认这是真竞态，不是理论风险。
	// 未加锁时的生产表现是偶发 nil 解引用崩溃（读到半个接口值）。
	//
	// 约定：只在持锁期间取字段值，取完立刻释放再调用。
	// 持锁调用会把 InjectInputSync 这类阻塞到 agent 回复（可达数分钟）的
	// 方法与 SetIOInjector 串到一起，让插件重载卡死。
	apiMu sync.RWMutex

	autoRestart bool

	stopMu       sync.Mutex
	stopHandlers []func()

	removeMu       sync.Mutex
	removeHandlers []func()
}

// New creates a PluginSDK with the given dependencies.
func New(name string, sett SettingsAPI, regTool ToolRegistrar, regStage StageRegistrar, regAPI APIRegistrar, regOutput OutputChannelRegistrar) *PluginSDK {
	return &PluginSDK{
		name:        name,
		sett:        sett,
		regTool:     regTool,
		regStage:    regStage,
		regAPI:      regAPI,
		regOutput:   regOutput,
		autoRestart: true,
	}
}

// PluginName returns the name of the plugin.
func (s *PluginSDK) PluginName() string { return s.name }

// Settings returns the settings API for reading/writing plugin configuration.
// sett 在 New 时一次性写入且无 setter，故不需要加锁。
func (s *PluginSDK) Settings() SettingsAPI { return s.sett }

// Memory returns the graph memory API (may be nil if not available).
func (s *PluginSDK) Memory() MemoryAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.mem
}

// TextMemory returns the text memory API (may be nil if not available).
func (s *PluginSDK) TextMemory() TextMemoryAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.textMem
}

// DocMemory returns the document memory API (may be nil if not available).
func (s *PluginSDK) DocMemory() DocMemoryAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.docMem
}

// Knowledge returns the knowledge store API (may be nil if not available).
func (s *PluginSDK) Knowledge() KnowledgeAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.know
}

// LLM returns the LLM provider API (may be nil if not available).
func (s *PluginSDK) LLM() LLMAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.llm
}

// Social returns the social graph API (may be nil if not available).
func (s *PluginSDK) Social() SocialAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.social
}

// Events returns the event subscriber for listening to kernel events (may be nil if not available).
func (s *PluginSDK) Events() EventSubscriber {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.events
}

// RegisterTool registers a tool that the LLM can call.
func (s *PluginSDK) RegisterTool(name string, def ToolDef, handler ToolHandler) error {
	if def.Plugin == "" {
		def.Plugin = s.name
	}
	if s.regTool != nil {
		return s.regTool(name, def, handler)
	}
	return nil
}

// RegisterStage registers a handler for a pipeline stage.
//
//	scope: StageScopeGlobal (default) — receives all stage events.
//	       StageScopeOwnTools — only before_toolcall/after_toolcall for this plugin's tools.
func (s *PluginSDK) RegisterStage(stage Stage, handler StageHandler, scope ...StageScope) {
	if s.regStage == nil {
		return
	}
	sc := StageScopeGlobal
	if len(scope) > 0 {
		sc = scope[0]
	}
	if sc == StageScopeGlobal {
		s.regStage(stage, handler)
		return
	}
	// OwnTools scope — only for before_toolcall / after_toolcall
	if stage != StageBeforeToolcall && stage != StageAfterToolcall {
		s.regStage(stage, handler)
		return
	}
	s.regStage(stage, func(ctx *StageContext) error {
		ctx.RLock()
		match := false
		switch stage {
		case StageBeforeToolcall:
			match = len(ctx.ToolCalls) > 0 && ctx.ToolCalls[0].Plugin == s.name
		case StageAfterToolcall:
			match = len(ctx.ToolResults) > 0 && ctx.ToolResults[0].Plugin == s.name
		}
		ctx.RUnlock()
		if !match {
			return nil
		}
		return handler(ctx)
	})
}

// RegisterPluginAPI registers this plugin's API for access by other plugins.
func (s *PluginSDK) RegisterPluginAPI(name string) error {
	if s.regAPI != nil {
		return s.regAPI(name)
	}
	return nil
}

// RegisterOutputChannel registers an output channel that the output_send tool can route to.
// name: channel name (e.g. "qq", "webui")
// caps: bitmask of supported output capabilities (CapText, CapFile, etc.)
// desc: description of the channel, expected meta format, and type enum
// def:  通道在记忆计算层的行为（NoMemory/Cleaner）
// handler: receives args map with keys: payload (string), type (string), meta (string|optional)
func (s *PluginSDK) RegisterOutputChannel(name string, caps int, desc string, def ChannelDef, handler ToolHandler) error {
	s.apiMu.RLock()
	reg := s.regOutput
	s.apiMu.RUnlock()
	if reg != nil {
		return reg(name, caps, desc, def, handler)
	}
	return nil
}

// RegisterInputChannel registers an input channel with its memory behavior.
// def.NoMemory: 此通道输入不参与记忆计算
// def.Cleaner:  计算层对输入文本清洗后（不改原文）再向量化/提关键词
func (s *PluginSDK) RegisterInputChannel(name string, def ChannelDef) error {
	s.apiMu.RLock()
	reg := s.regInput
	s.apiMu.RUnlock()
	if reg != nil {
		return reg(name, def)
	}
	return nil
}

// 以下 setter 由内核在启动/重载时调用，与插件后台 goroutine 的读并发，故加锁。

// SetOutputChannelRegistrar sets the output channel registrar (called by the core at startup).
func (s *PluginSDK) SetOutputChannelRegistrar(r OutputChannelRegistrar) {
	s.apiMu.Lock()
	s.regOutput = r
	s.apiMu.Unlock()
}

// SetInputChannelRegistrar sets the input channel registrar (called by the core at startup).
func (s *PluginSDK) SetInputChannelRegistrar(r InputChannelRegistrar) {
	s.apiMu.Lock()
	s.regInput = r
	s.apiMu.Unlock()
}

// SetIOInjector sets the IO injector (called by the core at startup).
func (s *PluginSDK) SetIOInjector(io IOInjector) {
	s.apiMu.Lock()
	s.io = io
	s.apiMu.Unlock()
}

// SetMemoryAPI sets the memory API (called by the core at startup).
func (s *PluginSDK) SetMemoryAPI(mem MemoryAPI) {
	s.apiMu.Lock()
	s.mem = mem
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetTextMemoryAPI(tm TextMemoryAPI) {
	s.apiMu.Lock()
	s.textMem = tm
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetDocMemoryAPI(dm DocMemoryAPI) {
	s.apiMu.Lock()
	s.docMem = dm
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetKnowledgeAPI(kn KnowledgeAPI) {
	s.apiMu.Lock()
	s.know = kn
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetLLMAPI(llm LLMAPI) {
	s.apiMu.Lock()
	s.llm = llm
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetSocialAPI(social SocialAPI) {
	s.apiMu.Lock()
	s.social = social
	s.apiMu.Unlock()
}

func (s *PluginSDK) SetEventSubscriber(es EventSubscriber) {
	s.apiMu.Lock()
	s.events = es
	s.apiMu.Unlock()
}

// SetPluginMgrAPI sets the plugin manager API (called by the bridge at startup).
func (s *PluginSDK) SetPluginMgrAPI(pm PluginMgrAPI) {
	s.apiMu.Lock()
	s.plgMgr = pm
	s.apiMu.Unlock()
}

// PluginMgr returns the plugin manager API (ReloadOne / ReloadPlugins / list).
// May be nil if the host did not wire it.
func (s *PluginSDK) PluginMgr() PluginMgrAPI {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.plgMgr
}

// ---- IO Convenience Methods ----

// injector 取当前 injector 的快照。
//
// 取完即释放锁再调用：InjectInputSync 会阻塞到 agent 回复（可达数分钟），
// 若持锁调用，插件重载时的 SetIOInjector 会一起卡住。
func (s *PluginSDK) injector() IOInjector {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.io
}

// InjectInterruptText injects a text interrupt that can preempt current LLM processing.
func (s *PluginSDK) InjectInterruptText(source, channel, text string) {
	if io := s.injector(); io != nil {
		io.InjectInterruptText(source, channel, text)
	}
}

// InjectText injects a text message into the agent pipeline.
func (s *PluginSDK) InjectText(source, channel, text string) {
	if io := s.injector(); io != nil {
		io.InjectText(source, channel, text)
	}
}

// InjectTextNoMemory injects a text message without generating memory.
func (s *PluginSDK) InjectTextNoMemory(source, channel, text string) {
	if io := s.injector(); io != nil {
		io.InjectTextNoMemory(source, channel, text)
	}
}

// InjectInputSync injects a text message and synchronously waits for the agent reply,
// returning the reply text (empty string if none). Replies must be dispatched back
// to the source channel by the caller.
func (s *PluginSDK) InjectInputSync(source, channel, text string) string {
	io := s.injector()
	if io == nil {
		return ""
	}
	return io.InjectInputSync(source, channel, text)
}

// InjectInputMedia 注入带媒体内容块（image_url/audio_url）的输入。
// blocks 会落进媒体存储被记忆引用捕获，同时作为当前轮 content 数组
// 发给 LLM，让模型在「本轮」就看到图/听到音频——区别于 SetToolBlocks
// 的「下一轮 tool message」语义。
func (s *PluginSDK) InjectInputMedia(source, channel, text string, blocks []ContentBlock) {
	if io := s.injector(); io != nil {
		io.InjectInputMedia(source, channel, text, blocks)
	}
}

// InjectInputMediaSync 注入带媒体内容块的输入并同步等待 agent 回复。
func (s *PluginSDK) InjectInputMediaSync(source, channel, text string, blocks []ContentBlock) string {
	io := s.injector()
	if io == nil {
		return ""
	}
	return io.InjectInputMediaSync(source, channel, text, blocks)
}

// InjectInterruptMedia 注入带媒体内容块的中断，可抢占当前 LLM 处理。
// blocks 随中断消息一起发给模型。
func (s *PluginSDK) InjectInterruptMedia(source, channel, text string, blocks []ContentBlock) {
	if io := s.injector(); io != nil {
		io.InjectInterruptMedia(source, channel, text, blocks)
	}
}

// SetToolBlocks 在工具处理函数内注入多模态内容块，内核在下一条 tool message
// 的 content 数组里带上它们。需要「本轮就让模型看到」时用 InjectInputMedia。
func (s *PluginSDK) SetToolBlocks(blocks []ContentBlock) {
	if io := s.injector(); io != nil {
		io.SetToolBlocks(blocks)
	}
}

// SetAutoRestart 设置插件是否允许内核自动重启（崩溃后自动重载）。
// 默认 true。如果插件有无法恢复的状态（如外部连接），应设为 false。
func (s *PluginSDK) SetAutoRestart(enabled bool) {
	s.apiMu.Lock()
	s.autoRestart = enabled
	s.apiMu.Unlock()
}

// AutoRestart 返回插件是否允许自动重启。
func (s *PluginSDK) AutoRestart() bool {
	s.apiMu.RLock()
	defer s.apiMu.RUnlock()
	return s.autoRestart
}

// RegisterStopHandler 注册插件停止阶段的清理回调。
// 注册的 handler 会在插件 Stop() 之前按"后注册先执行"的顺序调用，
// 适用于释放资源、落盘状态、关闭子进程等停止时清理操作。
// 可注册多个；执行后清空（进程停止前只执行一次）。
func (s *PluginSDK) RegisterStopHandler(fn func()) {
	if fn == nil {
		return
	}
	s.stopMu.Lock()
	s.stopHandlers = append(s.stopHandlers, fn)
	s.stopMu.Unlock()
}

// RunStopHandlers 执行全部已注册的 stop handler（后注册先执行，执行后清空，幂等）。
// 由内核（内置插件）或插件桥接层（外部插件 z_bridge 的 StopPlugin）在调用插件 Stop() 前执行。
func (s *PluginSDK) RunStopHandlers() {
	s.stopMu.Lock()
	handlers := append([]func(){}, s.stopHandlers...)
	s.stopHandlers = nil
	s.stopMu.Unlock()
	for i := len(handlers) - 1; i >= 0; i-- {
		handlers[i]()
	}
}

// RegisterOnRemoveHandler 注册插件被删除（卸载）时的清理回调。
// 注册的 handler 会在插件目录被移除前按"后注册先执行"的顺序调用，
// 适用于清理外部资源、删除配置表、下线状态等删除后处理。
// 可注册多个；执行后清空（一次删除只执行一次）。
func (s *PluginSDK) RegisterOnRemoveHandler(fn func()) {
	if fn == nil {
		return
	}
	s.removeMu.Lock()
	s.removeHandlers = append(s.removeHandlers, fn)
	s.removeMu.Unlock()
}

// RunOnRemoveHandlers 执行全部已注册的 onRemove handler（后注册先执行，执行后清空，幂等）。
// 由内核在卸载插件（registry.RemovePlugin）时、插件 Stop() 之后执行。
func (s *PluginSDK) RunOnRemoveHandlers() {
	s.removeMu.Lock()
	handlers := append([]func(){}, s.removeHandlers...)
	s.removeHandlers = nil
	s.removeMu.Unlock()
	for i := len(handlers) - 1; i >= 0; i-- {
		handlers[i]()
	}
}

// ContentBlock 是多模态内容块（OpenAI 格式：text/image_url/audio_url）。
// 插件工具返回结果时可用 PluginSDK.SetToolBlocks 注入，让下一轮 LLM
// 请求在 tool message 的 content 数组里带上图片/音频，实现"模型看图/听音频"。
type ContentBlock struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageURL *ImageURL `json:"image_url,omitempty"`
	AudioURL *AudioURL `json:"audio_url,omitempty"`
}

type ImageURL struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"`
}

type AudioURL struct {
	URL string `json:"url"`
}