From 249faf02b1518d795d608bd6c0b18f45919ced9b Mon Sep 17 00:00:00 2001
From: Dongbo Wang <dongbow@microsoft.com>
Date: Mon, 22 Jan 2024 18:09:38 -0800
Subject: [PATCH] Use separate `AssemblyLoadContext` for every agent plugin
 (#56)

---
 shell/ShellCopilot.Kernel/LLMAgent.cs         |  4 ++-
 shell/ShellCopilot.Kernel/Shell.cs            | 10 +++---
 .../ShellCopilot.Kernel/Utility/Clipboard.cs  |  3 +-
 .../Utility/LoadContext.cs                    | 36 +++++++++++++++++++
 4 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 shell/ShellCopilot.Kernel/Utility/LoadContext.cs

diff --git a/shell/ShellCopilot.Kernel/LLMAgent.cs b/shell/ShellCopilot.Kernel/LLMAgent.cs
index 2fbcf697..d5de1dd4 100644
--- a/shell/ShellCopilot.Kernel/LLMAgent.cs
+++ b/shell/ShellCopilot.Kernel/LLMAgent.cs
@@ -6,12 +6,14 @@ namespace ShellCopilot.Kernel;
 internal class LLMAgent
 {
     internal ILLMAgent Impl { get; }
+    internal AgentAssemblyLoadContext LoadContext { get; }
     internal bool OrchestratorRoleDisabled { set; get; }
     internal bool AnalyzerRoleDisabled { set; get; }
 
-    internal LLMAgent(ILLMAgent agent)
+    internal LLMAgent(ILLMAgent agent, AgentAssemblyLoadContext loadContext)
     {
         Impl = agent;
+        LoadContext = loadContext;
 
         OrchestratorRoleDisabled = false;
         AnalyzerRoleDisabled = false;
diff --git a/shell/ShellCopilot.Kernel/Shell.cs b/shell/ShellCopilot.Kernel/Shell.cs
index 672ac3f7..bb39d6ef 100644
--- a/shell/ShellCopilot.Kernel/Shell.cs
+++ b/shell/ShellCopilot.Kernel/Shell.cs
@@ -96,9 +96,11 @@ internal List<string> GetCodeBlockFromLastResponse()
     /// <summary>
     /// Load a plugin assembly file and process the agents defined in it.
     /// </summary>
-    internal void ProcessAgentPlugin(string pluginFile)
+    internal void ProcessAgentPlugin(string pluginName, string pluginDir, string pluginFile)
     {
-        Assembly plugin = Assembly.LoadFrom(pluginFile);
+        AgentAssemblyLoadContext context = new(pluginName, pluginDir);
+        Assembly plugin = context.LoadFromAssemblyPath(pluginFile);
+
         foreach (Type type in plugin.ExportedTypes)
         {
             if (!typeof(ILLMAgent).IsAssignableFrom(type))
@@ -121,7 +123,7 @@ internal void ProcessAgentPlugin(string pluginFile)
             };
 
             agent.Initialize(config);
-            _agents.Add(new LLMAgent(agent));
+            _agents.Add(new LLMAgent(agent, context));
         }
     }
 
@@ -147,7 +149,7 @@ private void LoadAvailableAgents()
 
             try
             {
-                ProcessAgentPlugin(file);
+                ProcessAgentPlugin(name, dir, file);
             }
             catch (Exception ex)
             {
diff --git a/shell/ShellCopilot.Kernel/Utility/Clipboard.cs b/shell/ShellCopilot.Kernel/Utility/Clipboard.cs
index 2613805f..6124ff83 100644
--- a/shell/ShellCopilot.Kernel/Utility/Clipboard.cs
+++ b/shell/ShellCopilot.Kernel/Utility/Clipboard.cs
@@ -1,5 +1,4 @@
-using System;
-using System.Diagnostics;
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 using System.Runtime.Versioning;
 
diff --git a/shell/ShellCopilot.Kernel/Utility/LoadContext.cs b/shell/ShellCopilot.Kernel/Utility/LoadContext.cs
new file mode 100644
index 00000000..644c681e
--- /dev/null
+++ b/shell/ShellCopilot.Kernel/Utility/LoadContext.cs
@@ -0,0 +1,36 @@
+using System.Reflection;
+using System.Runtime.Loader;
+
+namespace ShellCopilot.Kernel;
+
+internal class AgentAssemblyLoadContext : AssemblyLoadContext
+{
+    private readonly string _dependencyDir;
+
+    internal AgentAssemblyLoadContext(string name, string dependencyDir)
+        : base($"{name.Replace(' ', '.')}-ALC", isCollectible: false)
+    {
+        if (!Directory.Exists(dependencyDir))
+        {
+            throw new ArgumentException($"The agent home directory '{dependencyDir}' doesn't exist.", nameof(dependencyDir));
+        }
+
+        // Save the full path to the dependencies directory when creating the context.
+        _dependencyDir = dependencyDir;
+    }
+
+    protected override Assembly Load(AssemblyName assemblyName)
+    {
+        // Create a path to the assembly in the dependencies directory.
+        string path = Path.Combine(_dependencyDir, $"{assemblyName.Name}.dll");
+
+        if (File.Exists(path))
+        {
+            // If the assembly exists in our dependency directory, then load it into this load context.
+            return LoadFromAssemblyPath(path);
+        }
+
+        // Otherwise we will depend on the default load context to resolve the request.
+        return null;
+    }
+}