fix: Ai judges to increase the cache mechanism

internal/plugin/api.go

@@ -320,7 +320,7 @@ func (s *Server) handlePluginBind(c *wkhttp.Context) {
 	}
 
 	s.removePluginNoFromCache(req.Uid)
-
+	s.removeIsAiFromCache(req.Uid)
 	c.ResponseOK()
 }
 
@@ -373,7 +373,7 @@ func (s *Server) handlePluginUnbind(c *wkhttp.Context) {
 	}
 
 	s.removePluginNoFromCache(req.Uid)
-
+	s.removeIsAiFromCache(req.Uid)
 	c.ResponseOK()
 }
 

internal/plugin/model.go

@@ -103,16 +103,18 @@ func newPluginUserResp(pu wkdb.PluginUser) *pluginUserResp {
 }
 
 type userPluginBucket struct {
-	cache *lru.Cache[string, string]
-	index int
+	cache     *lru.Cache[string, string]
+	index     int
+	isAiCache *lru.Cache[string, bool] // 缓存用户是否是AI
 }
 
 func newUserPluginBucket(index int) *userPluginBucket {
 	cache, _ := lru.New[string, string](1000)
-
+	isAiCache, _ := lru.New[string, bool](1000)
 	return &userPluginBucket{
-		cache: cache,
-		index: index,
+		cache:     cache,
+		index:     index,
+		isAiCache: isAiCache,
 	}
 }
 
@@ -131,3 +133,16 @@ func (u *userPluginBucket) get(key string) (string, bool) {
 func (u *userPluginBucket) remove(key string) {
 	u.cache.Remove(key)
 }
+
+func (u *userPluginBucket) isAi(key string) (bool, bool) {
+	v, ok := u.isAiCache.Get(key)
+	return v, ok
+}
+
+func (u *userPluginBucket) setIsAi(key string, value bool) {
+	u.isAiCache.Add(key, value)
+}
+
+func (u *userPluginBucket) removeIsAi(key string) {
+	u.isAiCache.Remove(key)
+}

internal/plugin/server.go

@@ -149,15 +149,16 @@ func (s *Server) Plugin(no string) types.Plugin {
 }
 
 func (s *Server) UserIsAI(uid string) bool {
-	_, ok := s.getPluginNoFromCache(uid)
+	isAi, ok := s.isAiFromCache(uid)
 	if ok {
-		return ok
+		return isAi
 	}
 	exist, err := service.Store.DB().ExistPluginByUid(uid)
 	if err != nil {
 		s.Error("查询用户AI插件失败！", zap.Error(err), zap.String("uid", uid))
 		return false
 	}
+	s.setIsAiToCache(uid, exist)
 	return exist
 }
 
@@ -195,6 +196,24 @@ func (s *Server) removePluginNoFromCache(uid string) {
 	s.userPluginBuckets[index].remove(uid)
 }
 
+func (s *Server) setIsAiToCache(uid string, isAi bool) {
+	fh := fasthash.Hash(uid)
+	index := int(fh) % s.bucketSize
+	s.userPluginBuckets[index].setIsAi(uid, isAi)
+}
+
+func (s *Server) isAiFromCache(uid string) (bool, bool) {
+	fh := fasthash.Hash(uid)
+	index := int(fh) % s.bucketSize
+	return s.userPluginBuckets[index].isAi(uid)
+}
+
+func (s *Server) removeIsAiFromCache(uid string) {
+	fh := fasthash.Hash(uid)
+	index := int(fh) % s.bucketSize
+	s.userPluginBuckets[index].removeIsAi(uid)
+}
+
 func getUnixSocket(socketPath string) (string, error) {
 
 	err := os.Remove(socketPath)