Fix agent client init issue and clean up

src/Cody.Core/Agent/Connector/AgentConnector.cs

@@ -3,9 +3,6 @@
 using Newtonsoft.Json.Serialization;
 using StreamJsonRpc;
 using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
 using System.Threading.Tasks;
 
 namespace Cody.Core.Agent.Connector
@@ -54,13 +51,16 @@ public async Task Connect()
 
             if (options.NotificationsTarget != null) jsonRpc.AddLocalRpcTarget(options.NotificationsTarget);
             agentClient = jsonRpc.Attach<IAgentClient>();
-            options.NotificationsTarget.SetAgentClient(agentClient);
 
             jsonRpc.StartListening();
             IsConnected = true;
-            log.Info("A connection with the agent has been established.");
 
-            if (options.AfterConnection != null) options.AfterConnection(agentClient);
+            if (options.AfterConnection != null) await options.AfterConnection(agentClient);
+
+            // Makes sure the notifications target is set after the connection is established.
+            options.NotificationsTarget.SetAgentClient(agentClient);
+
+            log.Info("A connection with the agent has been established.");
         }
 
         private void OnAgentExit(int exitCode)

src/Cody.Core/Agent/Connector/AgentConnectorOptions.cs

@@ -18,8 +18,7 @@ public class AgentConnectorOptions
         public int? Port { get; set; }
 
         public string AgentDirectory { get; set; }
-
-        public Action<IAgentClient> AfterConnection { get; set; }
+        public Func<IAgentClient, Task> AfterConnection { get; set; }
 
         public Action<IAgentClient> BeforeDisconnection { get; set; }
 

src/Cody.Core/Agent/Connector/InitializeCallback.cs

@@ -5,10 +5,7 @@
 using Cody.Core.Logging;
 using Cody.Core.Settings;
 using System;
-using System.Collections.Generic;
 using System.IO;
-using System.Linq;
-using System.Text;
 using System.Threading.Tasks;
 
 namespace Cody.Core.Agent.Connector
@@ -91,13 +88,6 @@ public async Task Initialize(IAgentClient client)
             }
 
             client.Initialized();
-            await client.ResolveWebviewView(new ResolveWebviewViewParams
-            {
-                // cody.chat for sidebar view, or cody.editorPanel for editor panel
-                ViewId = "cody.chat",
-                // TODO: Create dynmically when we support editor panel
-                WebviewHandle = "visual-studio-sidebar",
-            });
         }
     }
 }

src/Cody.Core/Agent/NotificationHandlers.cs

@@ -1,10 +1,8 @@
-using Cody.Core.Agent.Connector;
-using Cody.Core.Agent.Protocol;
+using Cody.Core.Agent.Protocol;
+using EnvDTE80;
 using Newtonsoft.Json.Linq;
 using StreamJsonRpc;
 using System;
-using EnvDTE;
-using EnvDTE80;
 using System.Threading.Tasks;
 
 namespace Cody.Core.Agent
@@ -22,7 +20,14 @@ public NotificationHandlers()
 
         public IAgentClient agentClient;
 
-        public void SetAgentClient(IAgentClient agentClient) => this.agentClient = agentClient;
+        private TaskCompletionSource<bool> agentClientReady = new TaskCompletionSource<bool>();
+
+
+        public void SetAgentClient(IAgentClient client)
+        {
+            this.agentClient = client;
+            agentClientReady.SetResult(true);
+        }
 
         // Send a message to the host from webview.
         public async Task SendWebviewMessage(string handle, string message)
@@ -85,14 +90,15 @@ public void RegisterWebview(string handle)
         [JsonRpcMethod("webview/registerWebviewViewProvider")]
         public async Task RegisterWebviewViewProvider(string viewId, bool retainContextWhenHidden)
         {
-            System.Diagnostics.Debug.WriteLine(viewId, retainContextWhenHidden, "Agent registerWebviewViewProvider");
+            agentClientReady.Task.Wait();
             await agentClient.ResolveWebviewView(new ResolveWebviewViewParams
             {
                 // cody.chat for sidebar view, or cody.editorPanel for editor panel
-                ViewId = "cody.chat",
+                ViewId = viewId,
                 // TODO: Create dynmically when we support editor panel
                 WebviewHandle = "visual-studio-sidebar",
             });
+            System.Diagnostics.Debug.WriteLine(viewId, retainContextWhenHidden, "Agent registerWebviewViewProvider");
         }
 
         [JsonRpcMethod("webview/createWebviewPanel", UseSingleObjectParameterDeserialization = true)]

src/Cody.UI/Controls/WebView2Dev.xaml

@@ -16,7 +16,7 @@
                 <RowDefinition Height="*" />
             </Grid.RowDefinitions>
 
-            <wpf:WebView2 Name="webView" Grid.Row="1"
+            <wpf:WebView2 Name="webView" Grid.Row="1" Loaded="InitWebView2"
 
             />
         </Grid>

src/Cody.UI/Controls/WebView2Dev.xaml.cs

@@ -1,11 +1,11 @@
-using System;
+using Cody.Core.Agent;
+using Microsoft.Web.WebView2.Core;
+using System;
+using System.IO;
 using System.Threading.Tasks;
 using System.Windows;
 using System.Windows.Controls;
-using Microsoft.Web.WebView2.Core;
-using System.IO;
 using System.Windows.Input;
-using Cody.Core.Agent;
 
 namespace Cody.UI.Controls
 {
@@ -17,15 +17,16 @@ public partial class WebView2Dev : UserControl
 
         private static WebviewController _controller = new WebviewController();
 
+
         public WebView2Dev()
         {
             InitializeComponent();
-            InitializeWebView();
+            System.Diagnostics.Debug.WriteLine("InitializeComponent", "WebView2Dev");
         }
 
-        private void InitWebView2(object sender, RoutedEventArgs e)
+        private async void InitWebView2(object sender, RoutedEventArgs e)
         {
-            InitializeWebView();
+            await InitializeWebView();
         }
 
         public static WebviewController InitializeController(string themeScript)
@@ -34,14 +35,15 @@ public static WebviewController InitializeController(string themeScript)
             return _controller;
         }
 
-        private async Task<CoreWebView2> InitializeWebView()
+        private async Task InitializeWebView()
         {
             var env = await CreateWebView2Environment();
             await webView.EnsureCoreWebView2Async(env);
 
             _controller.WebViewMessageReceived += (sender, message) => SendMessage?.Execute(message);
 
-            return await _controller.InitializeWebView(webView.CoreWebView2);
+            await _controller.InitializeWebView(webView.CoreWebView2);
+            System.Diagnostics.Debug.WriteLine("InitializeWebView", "WebView2Dev");
         }
 
         private async Task<CoreWebView2Environment> CreateWebView2Environment()

src/Cody.UI/Controls/WebviewController.cs

@@ -1,9 +1,9 @@
-using System;
-using System.Threading.Tasks;
-using System.Windows;
 using Microsoft.Web.WebView2.Core;
+using System;
 using System.IO;
 using System.Reflection;
+using System.Threading.Tasks;
+using System.Windows;
 
 namespace Cody.UI.Controls
 {

src/Cody.UI/ViewModels/MainViewModel.cs

@@ -20,7 +20,7 @@ public MainViewModel(NotificationHandlers notificationHandlers, ILog logger)
             NotificationHandlers.OnSetHtmlEvent += OnSetHtmlHandler;
             NotificationHandlers.OnPostMessageEvent += OnPostMessageHandler;
 
-            _logger.Debug("Initialized.");
+            _logger.Debug("MainViewModel Initialized.");
         }
 
         private void OnPostMessageHandler(object sender, AgentResponseEvent e)

src/Cody.VisualStudio/CodyPackage.cs

@@ -1,31 +1,31 @@
-using Microsoft.VisualStudio.Shell;
+using Cody.Core.Agent;
+using Cody.Core.Agent.Connector;
+using Cody.Core.DocumentSync;
+using Cody.Core.Ide;
+using Cody.Core.Inf;
+using Cody.Core.Infrastructure;
+using Cody.Core.Logging;
+using Cody.Core.Settings;
+using Cody.UI.Controls;
+using Cody.UI.Views;
+using Cody.VisualStudio.Inf;
+using Cody.VisualStudio.Services;
+using Microsoft.VisualStudio.ComponentModelHost;
+using Microsoft.VisualStudio.Editor;
+using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Shell.Interop;
 using System;
 using System.ComponentModel.Design;
+using System.IO;
 using System.Linq;
+using System.Reflection;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Windows;
 using System.Windows.Threading;
-using Cody.Core.Ide;
-using Cody.Core.Inf;
-using Cody.Core.Logging;
-using Cody.UI.Views;
-using Cody.UI.Controls;
-using Cody.VisualStudio.Inf;
-using Cody.VisualStudio.Services;
 using Task = System.Threading.Tasks.Task;
-using System.Reflection;
-using System.IO;
-using Cody.Core.Settings;
-using Cody.Core.Infrastructure;
-using Cody.Core.Agent.Connector;
-using Cody.Core.Agent;
-using Cody.Core.DocumentSync;
-using Microsoft.VisualStudio.Editor;
-using Microsoft.VisualStudio.ComponentModelHost;
-using Microsoft.VisualStudio.Shell.Interop;
 
 namespace Cody.VisualStudio
 {
@@ -103,6 +103,7 @@ private void InitializeServices()
             StatusbarService = new StatusbarService();
             InitializeService = new InitializeCallback(UserSettingsService, VersionService, VsVersionService, StatusbarService, Logger);
             ThemeService = new ThemeService(this);
+            NotificationHandlers = new NotificationHandlers();
 
             var runningDocumentTable = this.GetService<SVsRunningDocumentTable, IVsRunningDocumentTable>();
             var componentModel = this.GetService<SComponentModel, IComponentModel>();
@@ -164,7 +165,6 @@ private async Task InitializeAgent()
             {
                 var agentDir = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), "Agent");
 
-                NotificationHandlers = new NotificationHandlers();
                 // Set the env var to 3113 when running with local agent.
                 var portNumber = int.TryParse(Environment.GetEnvironmentVariable("CODY_VS_DEV_PORT"), out int port) ? port : (int?)null;
 
@@ -173,17 +173,17 @@ private async Task InitializeAgent()
                     NotificationsTarget = NotificationHandlers,
                     AgentDirectory = agentDir,
                     RestartAgentOnFailure = true,
-                    AfterConnection = (client) => InitializeService.Initialize(client),
+                    AfterConnection = async (client) => await InitializeService.Initialize(client),
                     Port = portNumber,
                 };
 
-                AgentConnector = new AgentConnector(options, Logger);
-                AgentClientFactory = new AgentClientFactory(AgentConnector);
-
                 WebView2Dev.InitializeController(ThemeService.GetThemingScript());
                 NotificationHandlers.PostWebMessageAsJson = WebView2Dev.PostWebMessageAsJson;
 
-                _ = Task.Run(() => AgentConnector.Connect())
+                AgentConnector = new AgentConnector(options, Logger);
+                AgentClientFactory = new AgentClientFactory(AgentConnector);
+
+                _ = Task.Run(async () => await AgentConnector.Connect())
                 .ContinueWith(x =>
                 {
                     var documentSyncCallback = new DocumentSyncCallback(AgentClientFactory, Logger);