hooks.go 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303
  1. package remote_manager
  2. import (
  3. "encoding/hex"
  4. "sync"
  5. "sync/atomic"
  6. "time"
  7. "github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager/basic_manager"
  8. "github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager/media_manager"
  9. "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
  10. "github.com/langgenius/dify-plugin-daemon/internal/utils/cache"
  11. "github.com/langgenius/dify-plugin-daemon/internal/utils/log"
  12. "github.com/langgenius/dify-plugin-daemon/internal/utils/parser"
  13. "github.com/langgenius/dify-plugin-daemon/internal/utils/stream"
  14. "github.com/panjf2000/gnet/v2"
  15. )
  16. var (
  17. _mode pluginRuntimeMode
  18. )
  19. type DifyServer struct {
  20. gnet.BuiltinEventEngine
  21. engine gnet.Engine
  22. mediaManager *media_manager.MediaManager
  23. // listening address
  24. addr string
  25. port uint16
  26. // enabled multicore
  27. multicore bool
  28. // event loop count
  29. num_loops int
  30. // read new connections
  31. response *stream.Stream[*RemotePluginRuntime]
  32. plugins map[int]*RemotePluginRuntime
  33. plugins_lock *sync.RWMutex
  34. shutdown_chan chan bool
  35. }
  36. func (s *DifyServer) OnBoot(c gnet.Engine) (action gnet.Action) {
  37. s.engine = c
  38. return gnet.None
  39. }
  40. func (s *DifyServer) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
  41. // new plugin connected
  42. c.SetContext(&codec{})
  43. runtime := &RemotePluginRuntime{
  44. BasicPluginRuntime: basic_manager.NewBasicPluginRuntime(
  45. s.mediaManager,
  46. ),
  47. conn: c,
  48. response: stream.NewStream[[]byte](512),
  49. callbacks: make(map[string][]func([]byte)),
  50. callbacks_lock: &sync.RWMutex{},
  51. shutdown_chan: make(chan bool),
  52. alive: true,
  53. }
  54. // store plugin runtime
  55. s.plugins_lock.Lock()
  56. s.plugins[c.Fd()] = runtime
  57. s.plugins_lock.Unlock()
  58. // start a timer to check if handshake is completed in 10 seconds
  59. time.AfterFunc(time.Second*10, func() {
  60. if !runtime.handshake {
  61. // close connection
  62. c.Close()
  63. }
  64. })
  65. // verified
  66. verified := true
  67. if verified {
  68. return nil, gnet.None
  69. }
  70. return nil, gnet.Close
  71. }
  72. func (s *DifyServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
  73. // plugin disconnected
  74. s.plugins_lock.Lock()
  75. plugin := s.plugins[c.Fd()]
  76. delete(s.plugins, c.Fd())
  77. s.plugins_lock.Unlock()
  78. // close plugin
  79. plugin.onDisconnected()
  80. // uninstall plugin
  81. if plugin.assets_transferred {
  82. if _mode != _PLUGIN_RUNTIME_MODE_CI {
  83. if err := plugin.Unregister(); err != nil {
  84. log.Error("unregister plugin failed, error: %v", err)
  85. }
  86. }
  87. }
  88. // send stopped event
  89. plugin.wait_chan_lock.Lock()
  90. for _, c := range plugin.wait_stopped_chan {
  91. select {
  92. case c <- true:
  93. default:
  94. }
  95. }
  96. plugin.wait_chan_lock.Unlock()
  97. return gnet.None
  98. }
  99. func (s *DifyServer) OnShutdown(c gnet.Engine) {
  100. close(s.shutdown_chan)
  101. }
  102. func (s *DifyServer) OnTraffic(c gnet.Conn) (action gnet.Action) {
  103. codec := c.Context().(*codec)
  104. messages, err := codec.Decode(c)
  105. if err != nil {
  106. return gnet.Close
  107. }
  108. // get plugin runtime
  109. s.plugins_lock.RLock()
  110. runtime, ok := s.plugins[c.Fd()]
  111. s.plugins_lock.RUnlock()
  112. if !ok {
  113. return gnet.Close
  114. }
  115. // handle messages
  116. for _, message := range messages {
  117. if len(message) == 0 {
  118. continue
  119. }
  120. s.onMessage(runtime, message)
  121. }
  122. return gnet.None
  123. }
  124. func (s *DifyServer) onMessage(runtime *RemotePluginRuntime, message []byte) {
  125. // handle message
  126. if runtime.handshake_failed {
  127. // do nothing if handshake has failed
  128. return
  129. }
  130. close := func(message []byte) {
  131. if atomic.CompareAndSwapInt32(&runtime.closed, 0, 1) {
  132. runtime.conn.Write(message)
  133. runtime.conn.Close()
  134. }
  135. }
  136. if !runtime.handshake {
  137. key := string(message)
  138. info, err := GetConnectionInfo(key)
  139. if err == cache.ErrNotFound {
  140. // close connection if handshake failed
  141. close([]byte("handshake failed, invalid key\n"))
  142. runtime.handshake_failed = true
  143. return
  144. } else if err != nil {
  145. // close connection if handshake failed
  146. close([]byte("internal error\n"))
  147. return
  148. }
  149. runtime.tenant_id = info.TenantId
  150. // handshake completed
  151. runtime.handshake = true
  152. } else if !runtime.registration_transferred {
  153. // process handle shake if not completed
  154. declaration, err := parser.UnmarshalJsonBytes[plugin_entities.PluginDeclaration](message)
  155. if err != nil {
  156. // close connection if handshake failed
  157. close([]byte("handshake failed, invalid plugin declaration\n"))
  158. return
  159. }
  160. runtime.Config = declaration
  161. // registration transferred
  162. runtime.registration_transferred = true
  163. } else if !runtime.tools_registration_transferred {
  164. tools, err := parser.UnmarshalJsonBytes2Slice[plugin_entities.ToolProviderDeclaration](message)
  165. if err != nil {
  166. log.Error("tools register failed, error: %v", err)
  167. close([]byte("tools register failed, invalid tools declaration\n"))
  168. return
  169. }
  170. runtime.tools_registration_transferred = true
  171. if len(tools) > 0 {
  172. declaration := runtime.Config
  173. declaration.Tool = &tools[0]
  174. runtime.Config = declaration
  175. }
  176. } else if !runtime.models_registration_transferred {
  177. models, err := parser.UnmarshalJsonBytes2Slice[plugin_entities.ModelProviderDeclaration](message)
  178. if err != nil {
  179. log.Error("models register failed, error: %v", err)
  180. close([]byte("models register failed, invalid models declaration\n"))
  181. return
  182. }
  183. runtime.models_registration_transferred = true
  184. if len(models) > 0 {
  185. declaration := runtime.Config
  186. declaration.Model = &models[0]
  187. runtime.Config = declaration
  188. }
  189. } else if !runtime.endpoints_registration_transferred {
  190. endpoints, err := parser.UnmarshalJsonBytes2Slice[plugin_entities.EndpointProviderDeclaration](message)
  191. if err != nil {
  192. log.Error("endpoints register failed, error: %v", err)
  193. close([]byte("endpoints register failed, invalid endpoints declaration\n"))
  194. return
  195. }
  196. runtime.endpoints_registration_transferred = true
  197. if len(endpoints) > 0 {
  198. declaration := runtime.Config
  199. declaration.Endpoint = &endpoints[0]
  200. runtime.Config = declaration
  201. }
  202. } else if !runtime.assets_transferred {
  203. assets, err := parser.UnmarshalJsonBytes2Slice[plugin_entities.RemoteAssetPayload](message)
  204. if err != nil {
  205. log.Error("assets register failed, error: %v", err)
  206. close([]byte("assets register failed, invalid assets declaration\n"))
  207. return
  208. }
  209. files := make(map[string][]byte)
  210. for _, asset := range assets {
  211. files[asset.Filename], err = hex.DecodeString(asset.Data)
  212. if err != nil {
  213. log.Error("assets decode failed, error: %v", err)
  214. close([]byte("assets decode failed, invalid assets data, cannot decode file\n"))
  215. return
  216. }
  217. }
  218. // remap assets
  219. if err := runtime.RemapAssets(&runtime.Config, files); err != nil {
  220. log.Error("assets remap failed, error: %v", err)
  221. close([]byte("assets remap failed, invalid assets data, cannot remap\n"))
  222. return
  223. }
  224. runtime.assets_transferred = true
  225. runtime.checksum = runtime.calculateChecksum()
  226. runtime.InitState()
  227. runtime.SetActiveAt(time.Now())
  228. // trigger registration event
  229. if err := runtime.Register(); err != nil {
  230. log.Error("register failed, error: %v", err)
  231. close([]byte("register failed, cannot register\n"))
  232. return
  233. }
  234. // send started event
  235. runtime.wait_chan_lock.Lock()
  236. for _, c := range runtime.wait_started_chan {
  237. select {
  238. case c <- true:
  239. default:
  240. }
  241. }
  242. runtime.wait_chan_lock.Unlock()
  243. // publish runtime to watcher
  244. s.response.Write(runtime)
  245. } else {
  246. // continue handle messages if handshake completed
  247. runtime.response.Write(message)
  248. }
  249. }