GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/TypeAgent

Public

mirrored fromhttps://github.com/microsoft/TypeAgentAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
78c691a1dc6954d7bec8ba6b3df51daf1480fffe

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

TypeAgent/dotnet/visualStudioTypeAgent/Bridge

dotnet/visualStudioTypeAgent/Bridge/AgentBridgeClient.cs

207lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4using System;
5using System.Diagnostics;
6using System.Net.WebSockets;
7using System.Text;
8using System.Threading;
9using System.Threading.Tasks;
10using Microsoft.VisualStudio.Shell;
11using Newtonsoft.Json;
12using Newtonsoft.Json.Linq;
13
14namespace Microsoft.TypeAgent.VisualStudio.Bridge;
15
16/// <summary>
17/// WebSocket client that connects to the visualstudio-agent's bridge.
18/// Receives BridgeRequest messages, dispatches them through
19/// DTEActionExecutor, and sends BridgeResponse messages back.
20///
21/// Port discovery:
22/// The bridge port is no longer hardcoded. On each connect attempt
23/// we ask the agent-server's discovery channel where the
24/// `(visualStudio, default)` allocation lives. If discovery is
25/// unreachable or the agent isn't yet registered, the reconnect
26/// loop simply retries — there is no silent fallback to a
27/// well-known port. To pin a specific port (e.g. when running the
28/// bridge against a manually-launched agent), set
29/// `TYPEAGENT_VS_BRIDGE_PORT`; that bypasses discovery entirely.
30/// See <see cref="BridgeDiscovery"/> for the wire protocol and the
31/// `AGENT_SERVER_PORT` env-var knob.
32///
33/// Wire format (matches packages/agents/visualStudio/src/visualStudioActionHandler.ts):
34/// request: { id, actionName, parameters }
35/// response: { id, success, result?, error? }
36/// </summary>
37internal sealed class AgentBridgeClient : IDisposable
38{
39 private const string BridgePortOverrideEnv = "TYPEAGENT_VS_BRIDGE_PORT";
40 private static readonly TimeSpan ReconnectDelay = TimeSpan.FromSeconds(3);
41
42 private readonly AsyncPackage _package;
43 private readonly DTEActionExecutor _executor;
44 private readonly CancellationTokenSource _cts = new CancellationTokenSource();
45 private ClientWebSocket? _ws;
46
47 public AgentBridgeClient(AsyncPackage package)
48 {
49 _package = package;
50 _executor = new DTEActionExecutor(package);
51 }
52
53 public async Task StartAsync(CancellationToken cancellation)
54 {
55 using var linked = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, cancellation);
56 while (!linked.IsCancellationRequested)
57 {
58 int port = 0;
59 try
60 {
61 // Resolve the port fresh on every attempt: the agent may
62 // have restarted on a different ephemeral port since the
63 // last loop iteration, and the standalone shell may have
64 // come up while we were retrying.
65 uint? resolved = ResolvePortOverride()
66 ?? await BridgeDiscovery.ResolveBridgePortAsync(linked.Token).ConfigureAwait(false);
67 if (resolved is null)
68 {
69 // Discovery succeeded but the agent isn't registered
70 // yet — wait one reconnect cycle and try again.
71 Debug.WriteLine("[TypeAgent] visualStudio agent not yet registered; will retry");
72 }
73 else
74 {
75 port = (int)resolved.Value;
76 await ConnectAndReceiveAsync(port, linked.Token).ConfigureAwait(false);
77 }
78 }
79 catch (OperationCanceledException)
80 {
81 return;
82 }
83 catch (Exception ex)
84 {
85 Debug.WriteLine($"[TypeAgent] Bridge error (port {port}): {ex.Message}");
86 }
87 try
88 {
89 await Task.Delay(ReconnectDelay, linked.Token).ConfigureAwait(false);
90 }
91 catch (OperationCanceledException)
92 {
93 return;
94 }
95 }
96 }
97
98 // Returns an explicit port override from `TYPEAGENT_VS_BRIDGE_PORT`,
99 // or null when the env var is unset/malformed (caller falls through
100 // to discovery). Mirrors `CODE_WEBSOCKET_HOST` from coda.
101 private static uint? ResolvePortOverride()
102 {
103 string? raw = Environment.GetEnvironmentVariable(BridgePortOverrideEnv);
104 if (string.IsNullOrEmpty(raw))
105 {
106 return null;
107 }
108 if (uint.TryParse(raw, out uint p) && p > 0 && p <= 65535)
109 {
110 Debug.WriteLine($"[TypeAgent] {BridgePortOverrideEnv} override active: {p}");
111 return p;
112 }
113 Debug.WriteLine($"[TypeAgent] Ignoring malformed {BridgePortOverrideEnv}={raw}");
114 return null;
115 }
116
117 private async Task ConnectAndReceiveAsync(int port, CancellationToken cancellation)
118 {
119 var uri = new Uri($"ws://localhost:{port}");
120 _ws = new ClientWebSocket();
121 await _ws.ConnectAsync(uri, cancellation).ConfigureAwait(false);
122 Debug.WriteLine($"[TypeAgent] Bridge connected to {uri}");
123
124 var buffer = new ArraySegment<byte>(new byte[16 * 1024]);
125 var assembly = new StringBuilder();
126
127 while (_ws.State == WebSocketState.Open && !cancellation.IsCancellationRequested)
128 {
129 assembly.Clear();
130 WebSocketReceiveResult result;
131 do
132 {
133 result = await _ws.ReceiveAsync(buffer, cancellation).ConfigureAwait(false);
134 if (result.MessageType == WebSocketMessageType.Close)
135 {
136 await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cancellation).ConfigureAwait(false);
137 return;
138 }
139 assembly.Append(Encoding.UTF8.GetString(buffer.Array!, 0, result.Count));
140 } while (!result.EndOfMessage);
141
142 _ = HandleRequestAsync(assembly.ToString(), cancellation);
143 }
144 }
145
146 private async Task HandleRequestAsync(string json, CancellationToken cancellation)
147 {
148 string id = "";
149 try
150 {
151 var root = JObject.Parse(json);
152 id = root.Value<string>("id") ?? "";
153 var actionName = root.Value<string>("actionName") ?? "";
154 var parameters = root["parameters"] as JObject ?? new JObject();
155
156 var result = await _executor.ExecuteAsync(actionName, parameters, cancellation);
157 await SendResponseAsync(id, success: true, result: result, error: null, cancellation);
158 }
159 catch (Exception ex)
160 {
161 await SendResponseAsync(id, success: false, result: null, error: ex.Message, cancellation);
162 }
163 }
164
165 private async Task SendResponseAsync(string id, bool success, object? result, string? error, CancellationToken cancellation)
166 {
167 var ws = _ws;
168 if (ws is null || ws.State != WebSocketState.Open)
169 {
170 return;
171 }
172
173 var payload = new BridgeResponse
174 {
175 id = id,
176 success = success,
177 result = result,
178 error = error,
179 };
180 var json = JsonConvert.SerializeObject(payload, BridgeJson.Settings);
181 var bytes = Encoding.UTF8.GetBytes(json);
182 await ws.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Text, endOfMessage: true, cancellation);
183 }
184
185 public void Dispose()
186 {
187 _cts.Cancel();
188 try { _ws?.Dispose(); } catch { }
189 }
190
191 private sealed class BridgeResponse
192 {
193 public string id { get; set; } = "";
194 public bool success { get; set; }
195 public object? result { get; set; }
196 public string? error { get; set; }
197 }
198}
199
200internal static class BridgeJson
201{
202 // Preserve property names as authored (already camelCase) and drop nulls.
203 public static readonly JsonSerializerSettings Settings = new JsonSerializerSettings
204 {
205 NullValueHandling = NullValueHandling.Ignore,
206 };
207}
208