From 7f6845b2686b5a431f361dd17947afc85ee8ea11 Mon Sep 17 00:00:00 2001
From: "songguo.zeng" <songguo.zeng@starlight-sms.com>
Date: Thu, 12 Dec 2024 16:02:09 +0800
Subject: [PATCH] AII-400

---
 .../Agents/AgentHookBase.cs                   |  5 ++++
 .../BotSharp.Abstraction/Agents/IAgentHook.cs |  2 ++
 .../Agents/Services/AgentService.LoadAgent.cs | 27 ++++++++++++++++---
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/AgentHookBase.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/AgentHookBase.cs
index a3746ca8..f4d30523 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Agents/AgentHookBase.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Agents/AgentHookBase.cs
@@ -56,4 +56,9 @@ public abstract class AgentHookBase : IAgentHook
     public virtual void OnAgentUtilityLoaded(Agent agent)
     {
     }
+
+    public virtual void OnAgentLoadFilter(Agent agent)
+    {
+
+    }
 }
diff --git a/src/Infrastructure/BotSharp.Abstraction/Agents/IAgentHook.cs b/src/Infrastructure/BotSharp.Abstraction/Agents/IAgentHook.cs
index a3f53fb6..a329134b 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Agents/IAgentHook.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Agents/IAgentHook.cs
@@ -33,4 +33,6 @@ public interface IAgentHook
     /// <param name="agent"></param>
     /// <returns></returns>
     void OnAgentLoaded(Agent agent);
+
+    void OnAgentLoadFilter(Agent agent);
 }
diff --git a/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.LoadAgent.cs b/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.LoadAgent.cs
index 396a7d15..3ae49385 100644
--- a/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.LoadAgent.cs
+++ b/src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.LoadAgent.cs
@@ -5,9 +5,8 @@ namespace BotSharp.Core.Agents.Services;
 
 public partial class AgentService
 {
-    public static ConcurrentDictionary<string, Dictionary<string,string>> AgentParameterTypes = new();
+    public static ConcurrentDictionary<string, Dictionary<string, string>> AgentParameterTypes = new();
 
-    [MemoryCache(10 * 60, perInstanceCache: true)]
     public async Task<Agent> LoadAgent(string id)
     {
         if (string.IsNullOrEmpty(id) || id == Guid.Empty.ToString())
@@ -15,6 +14,26 @@ public partial class AgentService
             return null;
         }
 
+        Agent agent = await GetLoadAgent(id);
+        OnAgentLoadFilter(agent);
+        return agent;
+    }
+
+    private void OnAgentLoadFilter(Agent? agent)
+    {
+        if (agent != null && agent.Type == AgentType.Routing)
+        {
+            var hooks = _services.GetServices<IAgentHook>();
+            foreach (var hook in hooks)
+            {
+                hook.OnAgentLoadFilter(agent);
+            }
+        }
+    }
+
+    [MemoryCache(10 * 60, perInstanceCache: true)]
+    private async Task<Agent> GetLoadAgent(string id)
+    {
         var hooks = _services.GetServices<IAgentHook>();
 
         // Before agent is loaded.
@@ -106,14 +125,14 @@ public partial class AgentService
     {
         var agentId = agent.Id ?? agent.Name;
         if (AgentParameterTypes.ContainsKey(agentId)) return;
-        
+
         AddOrUpdateRoutesParameters(agentId, agent.RoutingRules);
         AddOrUpdateFunctionsParameters(agentId, agent.Functions);
     }
 
     private void AddOrUpdateRoutesParameters(string agentId, List<RoutingRule> routingRules)
     {
-        if(!AgentParameterTypes.TryGetValue(agentId, out var parameterTypes))
+        if (!AgentParameterTypes.TryGetValue(agentId, out var parameterTypes))
         {
             parameterTypes = new();
         }