Skip to content

Commit c61c4d5

Browse files
PiotrKarczmarzPiotrKarczmarz
authored
Unified Agent Lifecycle Management with AgentService (#376)
CLOSE https://linear.app/sourcegraph/issue/CODY-6258/visual-studio-plugin-crashes-on-disconnect Introduced a new `AgentService` class to centralize and unify all agent lifecycle management, ensuring proper initialization and restart workflows. Architectural Changes ### New AgentService Class (src/Cody.VisualStudio/Services/AgentService.cs) - Encapsulates complete agent initialization workflow (Start → Initialize → Component Setup) - Manages automatic restart with full re-initialization on disconnection - Provides centralized event handling for initialization success/failure - Eliminates duplicated initialization logic ### AgentClient Refactoring (src/Cody.VisualStudio/Client/AgentClient.cs) - Removed restart logic from OnAgentDisconnected() - Added AgentDisconnected event to delegate restart handling to AgentService - Agent is now focused solely on connection management, not lifecycle orchestration ### CodyPackage Integration (src/Cody.VisualStudio/CodyPackage.cs) - Replaced direct agent initialization with AgentServiceManager.InitializeAsync() - Added OnAgentServiceInitialized() callback for post-initialization component setup - Simplified InitializeAgent() method by delegating complexity to AgentService Benefits - ✅ Complete restart workflow: Agent now properly re-initializes after disconnection - ✅ Single responsibility: Clear separation between connection (AgentClient) and lifecycle (AgentService) - ✅ Maintainability: Centralized initialization logic easier to modify and test - ✅ Consistency: All agent startup paths go through the same workflow ## Test plan 1. Start VS Experimental instance 2. Kill Node.js process hosting the agent 3. Check "Cody" output window and follow the agent restart flow 4. Check the chat which should automatically refresh itself
1 parent 4d8788a commit c61c4d5

File tree

20 files changed

+787
-91
lines changed

20 files changed

+787
-91
lines changed

‎src/.contextkeeper/Agent Service Flow.ck‎

Lines changed: 487 additions & 0 deletions
Large diffs are not rendered by default.

src/Cody.Core/Agent/IAgentService.cs renamed to src/Cody.Core/Agent/IAgentApi.cs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
using Cody.Core.Agent.Protocol;
2-
using System;
3-
using System.Collections.Generic;
42
using System.Threading;
53
using System.Threading.Tasks;
64

@@ -9,7 +7,7 @@ namespace Cody.Core.Agent
97
//---------------------------------------------------------
108
// For notifications return type MUST be void!
119
//---------------------------------------------------------
12-
public interface IAgentService
10+
public interface IAgentApi
1311
{
1412
[AgentCall("initialize")]
1513
Task<ServerInfo> Initialize(ClientInfo clientInfo);

‎src/Cody.Core/Agent/IAgentProxy.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@ public interface IAgentProxy
99
bool IsConnected { get; }
1010
bool IsInitialized { get; }
1111
void Start();
12-
Task<IAgentService> Initialize(ClientInfo clientInfo);
12+
Task<IAgentApi> Initialize(ClientInfo clientInfo);
1313
event EventHandler<ServerInfo> OnInitialized;
14+
event EventHandler<int> AgentDisconnected;
1415
}
1516
}

‎src/Cody.Core/Agent/ProgressNotificationHandlers.cs‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,43 @@
11
using Cody.Core.Agent.Protocol;
22
using Cody.Core.Infrastructure;
3-
using Cody.Core.Logging;
43
using System;
5-
using System.Collections.Generic;
6-
using System.Linq;
7-
using System.Text;
8-
using System.Threading.Tasks;
94

105
namespace Cody.Core.Agent
116
{
127
public class ProgressNotificationHandlers
138
{
14-
private IProgressService progressService;
15-
private IAgentService agentService;
9+
private readonly IProgressService _progressService;
10+
private IAgentService _agentService;
1611

1712
public ProgressNotificationHandlers(IProgressService progressService)
1813
{
19-
this.progressService = progressService;
14+
this._progressService = progressService;
2015
}
2116

22-
public void SetAgentService(IAgentService agentService) => this.agentService = agentService;
17+
public void SetAgentService(IAgentService agentService) => this._agentService = agentService;
2318

2419
[AgentCallback("progress/start", deserializeToSingleObject: true)]
2520
public void Start(ProgressStartParams progressStart)
2621
{
2722
Action cancelAction = null;
2823
if (progressStart.Options.Cancellable == true)
2924
{
30-
cancelAction = () => agentService.CancelProgress(progressStart.Id);
25+
cancelAction = () => _agentService.Get().CancelProgress(progressStart.Id);
3126
};
3227

33-
progressService.Start(progressStart.Id, progressStart.Options.Title, cancelAction);
28+
_progressService.Start(progressStart.Id, progressStart.Options.Title, cancelAction);
3429
}
3530

3631
[AgentCallback("progress/report", deserializeToSingleObject: true)]
3732
public void Report(ProgressReportParams progressReport)
3833
{
39-
progressService.ReportProgress(progressReport.Id, progressReport.Message, progressReport.Increment);
34+
_progressService.ReportProgress(progressReport.Id, progressReport.Message, progressReport.Increment);
4035
}
4136

4237
[AgentCallback("progress/end")]
4338
public void End(string id)
4439
{
45-
progressService.End(id);
40+
_progressService.End(id);
4641
}
4742
}
4843
}

‎src/Cody.Core/Agent/WebviewNotificationHandlers.cs‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,15 @@
22
using Cody.Core.Logging;
33
using Newtonsoft.Json.Linq;
44
using System;
5-
using System.Collections.Generic;
6-
using System.Linq;
7-
using System.Text;
85
using System.Threading.Tasks;
6+
using Cody.Core.Infrastructure;
97

108
namespace Cody.Core.Agent
119
{
1210
public class WebviewNotificationHandlers
1311
{
1412
private readonly ILog _logger;
15-
public IAgentService agentClient;
13+
public IAgentService _agentService;
1614
private WebviewMessageHandler _messageFilter;
1715

1816
public WebviewNotificationHandlers(ILog logger)
@@ -26,21 +24,33 @@ public WebviewNotificationHandlers(ILog logger)
2624
public event EventHandler OnOptionsPageShowRequest;
2725
public event EventHandler<string> PostWebMessageAsJson;
2826

29-
public void SetAgentClient(IAgentService client)
27+
public void SetAgentClient(IAgentService agentService)
3028
{
31-
agentClient = client;
29+
_agentService = agentService;
3230
}
3331

3432
// Send a message to the host from webview.
3533
public async Task SendWebviewMessage(string handle, string message)
3634
{
37-
bool handled = _messageFilter.HandleMessage(message);
38-
if (!handled)
39-
await agentClient.ReceiveMessageStringEncoded(new ReceiveMessageStringEncodedParams
35+
try
36+
{
37+
bool handled = _messageFilter.HandleMessage(message);
38+
if (_agentService.Get() == null)
4039
{
41-
Id = handle,
42-
MessageStringEncoded = message
43-
});
40+
_logger.Error("There is no agent listening.");
41+
}
42+
43+
if (!handled && _agentService.Get() != null)
44+
await _agentService.Get().ReceiveMessageStringEncoded(new ReceiveMessageStringEncodedParams
45+
{
46+
Id = handle,
47+
MessageStringEncoded = message
48+
});
49+
}
50+
catch (Exception ex)
51+
{
52+
_logger.Error("Sending message to the agent failed.", ex);
53+
}
4454
}
4555

4656

‎src/Cody.Core/Cody.Core.csproj‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@
5050
<Compile Include="Agent\AgentCallAttribute.cs" />
5151
<Compile Include="Agent\AgentCallbackAttribute.cs" />
5252
<Compile Include="Agent\EditTaskNotificationHandlers.cs" />
53+
<Compile Include="Agent\IAgentApi.cs" />
5354
<Compile Include="Agent\IAgentProxy.cs" />
54-
<Compile Include="Agent\IAgentService.cs" />
5555
<Compile Include="Agent\NotificationHandlers.cs" />
5656
<Compile Include="Agent\ProgressNotificationHandlers.cs" />
5757
<Compile Include="Agent\Protocol\AutocompleteEditItem.cs" />
@@ -122,6 +122,7 @@
122122
<Compile Include="DocumentSync\IDocumentSyncActions.cs" />
123123
<Compile Include="Ide\IInfobarNotifications.cs" />
124124
<Compile Include="Ide\Notification.cs" />
125+
<Compile Include="Infrastructure\IAgentService.cs" />
125126
<Compile Include="Infrastructure\ConfigurationService.cs" />
126127
<Compile Include="Infrastructure\IConfigurationService.cs" />
127128
<Compile Include="Infrastructure\IDocumentService.cs" />

‎src/Cody.Core/DocumentSync/DocumentSyncCallback.cs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
using System;
77
using System.Collections.Generic;
88
using System.Linq;
9+
using Cody.Core.Infrastructure;
910

1011
namespace Cody.Core.DocumentSync
1112
{
@@ -80,7 +81,7 @@ public void OnChanged(string fullPath, DocumentRange visibleRange, DocumentRange
8081
}).ToArray()
8182
};
8283

83-
agentService.DidChange(docState);
84+
agentService.Get().DidChange(docState);
8485
}
8586

8687
public void OnClosed(string fullPath)
@@ -93,13 +94,13 @@ public void OnClosed(string fullPath)
9394
};
9495

9596
// Only the 'uri' property is required, other properties are ignored.
96-
agentService.DidClose(docState);
97+
agentService.Get().DidClose(docState);
9798
}
9899

99100
public void OnFocus(string fullPath)
100101
{
101102
trace.TraceEvent("DidFocus", "{0}", fullPath);
102-
agentService.DidFocus(new CodyFilePath { Uri = fullPath.ToUri() });
103+
agentService.Get().DidFocus(new CodyFilePath { Uri = fullPath.ToUri() });
103104
}
104105

105106
public void OnOpened(string fullPath, string content, DocumentRange visibleRange, DocumentRange selection)
@@ -152,7 +153,7 @@ public void OnOpened(string fullPath, string content, DocumentRange visibleRange
152153
Selection = selRange
153154
};
154155

155-
agentService.DidOpen(docState);
156+
agentService.Get().DidOpen(docState);
156157
}
157158
catch (Exception ex)
158159
{
@@ -163,13 +164,13 @@ public void OnOpened(string fullPath, string content, DocumentRange visibleRange
163164
public void OnRename(string oldFullPath, string newFullPath)
164165
{
165166
trace.TraceEvent("DidRename", "{0} -> {1}", oldFullPath, newFullPath);
166-
agentService.DidRename(new CodyFileRename { OldUri = oldFullPath.ToUri(), NewUri = newFullPath.ToUri() });
167+
agentService.Get().DidRename(new CodyFileRename { OldUri = oldFullPath.ToUri(), NewUri = newFullPath.ToUri() });
167168
}
168169

169170
public void OnSaved(string fullPath)
170171
{
171172
trace.TraceEvent("DidSave", "{0}", fullPath);
172-
agentService.DidSave(new CodyFilePath { Uri = fullPath.ToUri() });
173+
agentService.Get().DidSave(new CodyFilePath { Uri = fullPath.ToUri() });
173174
}
174175
}
175176
}

‎src/Cody.Core/Infrastructure/IAgentService.cs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
using Cody.Core.Agent;
2+
3+
namespace Cody.Core.Infrastructure
4+
{
5+
public interface IAgentService
6+
{
7+
IAgentApi Get();
8+
}
9+
}

‎src/Cody.Core/Infrastructure/WebViewsManager.cs‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,13 @@ public class WebViewsManager : IWebViewsManager, IDisposable
4242
private readonly ILog _logger;
4343

4444
private readonly List<IWebChatHost> _chatHosts;
45-
private readonly List<WebViewEvent> _processedWebViewsRequests;
45+
private ConcurrentQueue<WebViewEvent> _processedWebViewsRequests;
4646

4747
private readonly TimeSpan _agentInitializationTimeout = TimeSpan.FromMinutes(1);
4848

4949
private readonly BlockingCollection<WebViewEvent> _events; //TODO: when custom editors will be introduced, make it richer, like BlockingCollection<WebViewsEvents>, where WebViewsEvents will be a class
5050

51+
5152
public WebViewsManager(IAgentProxy agentProxy, WebviewNotificationHandlers notificationHandler, ILog logger)
5253
{
5354
_agentProxy = agentProxy;
@@ -56,8 +57,9 @@ public WebViewsManager(IAgentProxy agentProxy, WebviewNotificationHandlers notif
5657

5758
_chatHosts = new List<IWebChatHost>();
5859
_events = new BlockingCollection<WebViewEvent>();
59-
_processedWebViewsRequests = new List<WebViewEvent>();
60+
_processedWebViewsRequests = new ConcurrentQueue<WebViewEvent>();
6061

62+
_agentProxy.AgentDisconnected += OnAgentDisconnected;
6163
_notificationHandler.OnRegisterWebViewRequest += OnRegisterWebViewRequestHandler;
6264

6365
Task.Run(ProcessEvents);
@@ -79,11 +81,11 @@ private async Task ProcessEvents()
7981

8082
_logger.Debug($"Processing event '{e.Type}' ...");
8183

82-
_processedWebViewsRequests.Add(e);
84+
_processedWebViewsRequests.Enqueue(e);
8385
_logger.Debug($"{e.Type}:{e.ViewId} added to processed collection.");
8486

85-
var isRegisterWebViewRequestProcessed = _processedWebViewsRequests.Exists(w => w.Type == WebViewsEventTypes.RegisterWebViewRequest);
86-
var isWebChatHostInitialized = _processedWebViewsRequests.Exists(w => w.Type == WebViewsEventTypes.WebChatHostInitialized);
87+
var isRegisterWebViewRequestProcessed = _processedWebViewsRequests.Any(w => w.Type == WebViewsEventTypes.RegisterWebViewRequest);
88+
var isWebChatHostInitialized = _processedWebViewsRequests.Any(w => w.Type == WebViewsEventTypes.WebChatHostInitialized);
8789
if (isRegisterWebViewRequestProcessed && isWebChatHostInitialized)
8890
{
8991
// check if there is IWebChatHost available
@@ -99,7 +101,7 @@ private async Task ProcessEvents()
99101
try
100102
{
101103
await WaitForAgentInitialization();
102-
await _agentService.ResolveWebviewView(new ResolveWebviewViewParams
104+
await _agentService.Get().ResolveWebviewView(new ResolveWebviewViewParams
103105
{
104106
// cody.chat for sidebar view, or cody.editorPanel for editor panel
105107
// TODO support custom editors
@@ -129,10 +131,25 @@ await _agentService.ResolveWebviewView(new ResolveWebviewViewParams
129131
}
130132
}
131133

134+
private void OnAgentDisconnected(object sender, int e)
135+
{
136+
_logger.Debug("Cleaning ...");
137+
138+
var newQueue = new ConcurrentQueue<WebViewEvent>(
139+
_processedWebViewsRequests.Where(r => r.Type != WebViewsEventTypes.RegisterWebViewRequest)
140+
);
141+
142+
_processedWebViewsRequests = newQueue;
143+
144+
_logger.Debug("Cleared old agent service reference.");
145+
}
146+
132147
private async Task WaitForAgentInitialization()
133148
{
134149
var startTime = DateTime.Now;
135-
while (!_agentProxy.IsInitialized || _agentService == null)
150+
while (!_agentProxy.IsInitialized ||
151+
_agentService == null || _agentService.Get() == null
152+
)
136153
{
137154
_logger.Debug("Waiting for Agent initialization ...");
138155
await Task.Delay(TimeSpan.FromSeconds(1));
@@ -182,6 +199,12 @@ public void Register(IWebChatHost chatHost)
182199
public void Dispose()
183200
{
184201
_events?.Dispose();
202+
203+
if (_agentProxy != null)
204+
_agentProxy.AgentDisconnected -= OnAgentDisconnected;
205+
206+
if (_notificationHandler != null)
207+
_notificationHandler.OnRegisterWebViewRequest -= OnRegisterWebViewRequestHandler;
185208
}
186209
}
187210

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -114,15 +114,15 @@ public static async void PostWebMessageAsJson(object sender, string message)
114114
}
115115

116116
public static readonly DependencyProperty HtmlProperty = DependencyProperty.Register(
117-
"Html", typeof(string), typeof(WebView2Dev),
117+
"Html", typeof(SetHtmlEvent), typeof(WebView2Dev),
118118
new PropertyMetadata(null, (d, e) =>
119119
{
120-
_controller.SetHtml(e.NewValue as string);
120+
_controller.SetHtml(e.NewValue as SetHtmlEvent);
121121
}));
122122

123-
public string Html
123+
public SetHtmlEvent Html
124124
{
125-
get => (string)GetValue(HtmlProperty);
125+
get => (SetHtmlEvent)GetValue(HtmlProperty);
126126
set => SetValue(HtmlProperty, value);
127127
}
128128

0 commit comments

Comments
 (0)