GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/TypeAgent

Public

mirrored from https://github.com/microsoft/TypeAgentAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
copilot/fix-github-actions-job-shell-and-cli

Branches

Tags

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

Clone

HTTPS

Download ZIP

TypeAgent/dotnet/visualStudioTypeAgent/Bridge

dotnet/visualStudioTypeAgent/Bridge/BridgeDiscovery.cs

211lines · 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 Newtonsoft.Json;
11using Newtonsoft.Json.Linq;
12
13namespace Microsoft.TypeAgent.VisualStudio.Bridge;
14
15/// <summary>
16/// Looks up the visualStudio agent's bridge port via the dispatcher's
17/// discovery channel.
18///
19/// Wire protocol (matches packages/agentServer/protocol/src/protocol.ts):
20/// client → server:
21/// { "name": "discovery",
22/// "message": { "type": "invoke", "callId": N, "name": "lookupPort",
23/// "args": [{ "agentName": "visualStudio",
24/// "role": "default" }] } }
25/// server → client:
26/// { "name": "discovery",
27/// "message": { "type": "invokeResult", "callId": N,
28/// "result": { "port": &lt;int|null&gt; } } }
29///
30/// Returns the discovered port, or null when the agent isn't yet
31/// registered with the agent-server (transient — caller should retry).
32/// Throws on transport failure so the outer reconnect loop can apply
33/// its own retry/backoff. There is intentionally no hardcoded fallback
34/// port — the migrated TS clients (browser, code, coda) all return
35/// "undefined" on discovery failure and rely on the reconnect loop;
36/// dialing a stale well-known port would just connect to nothing.
37/// </summary>
38internal static class BridgeDiscovery
39{
40 // Read on every resolve so users can flip behavior without
41 // restarting the IDE between debugging sessions.
42 private const string AgentServerPortEnv = "AGENT_SERVER_PORT";
43
44 // Must match AGENT_SERVER_DEFAULT_PORT in agentServer/protocol.
45 private const uint DefaultAgentServerPort = 8999;
46
47 // Names this client uses to look itself up. Must match the role
48 // registered by visualStudioActionHandler.ts.
49 private const string AgentName = "visualStudio";
50 private const string Role = "default";
51
52 // The dispatcher's discovery channel always lives on the loopback
53 // agent-server. Keep this as a const so the URL only appears in
54 // one place; if/when the host becomes configurable, change here.
55 private const string AgentServerHost = "ws://localhost";
56
57 // Sanity cap on the discovery response payload. The protocol only
58 // ever returns a tiny JSON envelope (~100 bytes); anything larger
59 // is treated as a malformed/unexpected response.
60 private const int MaxDiscoveryResponseBytes = 64 * 1024;
61
62 /// <summary>
63 /// Resolve the bridge port via discovery. Returns the discovered
64 /// port, or null when the agent has not yet registered (transient
65 /// — caller should retry on its reconnect loop).
66 /// Throws on transport failure (agent-server unreachable, timeout,
67 /// malformed response) so the caller can log and retry.
68 /// </summary>
69 public static async Task<uint?> ResolveBridgePortAsync(CancellationToken cancellation)
70 {
71 uint agentServerPort = GetAgentServerPort();
72 uint? discovered = await LookupPortAsync(agentServerPort, cancellation).ConfigureAwait(false);
73 if (discovered is uint p)
74 {
75 Debug.WriteLine($"[TypeAgent] Discovery resolved bridge port {p}");
76 }
77 else
78 {
79 Debug.WriteLine($"[TypeAgent] Discovery returned null for ({AgentName}, {Role}); agent not yet registered");
80 }
81 return discovered;
82 }
83
84 private static uint GetAgentServerPort()
85 {
86 string? raw = Environment.GetEnvironmentVariable(AgentServerPortEnv);
87 if (uint.TryParse(raw, out uint p) && p > 0 && p <= 65535)
88 {
89 return p;
90 }
91 return DefaultAgentServerPort;
92 }
93
94 private static async Task<uint?> LookupPortAsync(uint agentServerPort, CancellationToken cancellation)
95 {
96 var uri = new Uri($"{AgentServerHost}:{agentServerPort}/");
97 using var ws = new ClientWebSocket();
98 // Cap the discovery call so a hung agent-server doesn't stall
99 // the whole reconnect loop. The outer AgentBridgeClient loop
100 // already retries on a separate cadence.
101 using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
102 using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellation, timeout.Token);
103
104 await ws.ConnectAsync(uri, linked.Token).ConfigureAwait(false);
105
106 // callId is arbitrary — the server echoes it back verbatim,
107 // and we only have one outstanding request per socket.
108 const int callId = 1;
109 var request = new JObject
110 {
111 ["name"] = "discovery",
112 ["message"] = new JObject
113 {
114 ["type"] = "invoke",
115 ["callId"] = callId,
116 ["name"] = "lookupPort",
117 ["args"] = new JArray
118 {
119 new JObject
120 {
121 ["agentName"] = AgentName,
122 ["role"] = Role,
123 },
124 },
125 },
126 };
127 byte[] requestBytes = Encoding.UTF8.GetBytes(request.ToString(Formatting.None));
128 await ws.SendAsync(
129 new ArraySegment<byte>(requestBytes),
130 WebSocketMessageType.Text,
131 endOfMessage: true,
132 linked.Token).ConfigureAwait(false);
133
134 string responseText = await ReceiveFullMessageAsync(ws, linked.Token).ConfigureAwait(false);
135 try
136 {
137 await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, CancellationToken.None).ConfigureAwait(false);
138 }
139 catch
140 {
141 // Best-effort close — the response is already in hand.
142 }
143
144 var root = JObject.Parse(responseText);
145 string? name = root.Value<string>("name");
146 if (name != "discovery")
147 {
148 return null;
149 }
150 var inner = root["message"] as JObject;
151 if (inner == null)
152 {
153 return null;
154 }
155 string? type = inner.Value<string>("type");
156 if (type == "invokeError")
157 {
158 throw new InvalidOperationException(
159 inner.Value<string>("error") ?? "Discovery returned invokeError");
160 }
161 if (type != "invokeResult")
162 {
163 return null;
164 }
165 if (inner.Value<int?>("callId") != callId)
166 {
167 return null;
168 }
169 var result = inner["result"] as JObject;
170 if (result == null)
171 {
172 return null;
173 }
174 // `port` is `int|null` on the wire; clamp to the valid port
175 // range and surface anything else as "not registered".
176 int? portValue = result.Value<int?>("port");
177 if (portValue is int pv && pv > 0 && pv <= 65535)
178 {
179 return (uint)pv;
180 }
181 return null;
182 }
183
184 private static async Task<string> ReceiveFullMessageAsync(ClientWebSocket ws, CancellationToken cancellation)
185 {
186 // 16KB receive chunk; we loop until EndOfMessage so this is a
187 // chunk size, not a hard message cap. The MaxDiscoveryResponseBytes
188 // guard below bounds the total payload to protect against a
189 // misbehaving peer streaming garbage.
190 var buffer = new ArraySegment<byte>(new byte[16 * 1024]);
191 var sb = new StringBuilder();
192 WebSocketReceiveResult result;
193 int totalBytes = 0;
194 do
195 {
196 result = await ws.ReceiveAsync(buffer, cancellation).ConfigureAwait(false);
197 if (result.MessageType == WebSocketMessageType.Close)
198 {
199 throw new InvalidOperationException("Discovery WS closed before response");
200 }
201 totalBytes += result.Count;
202 if (totalBytes > MaxDiscoveryResponseBytes)
203 {
204 throw new InvalidOperationException(
205 $"Discovery response exceeded {MaxDiscoveryResponseBytes} bytes; aborting");
206 }
207 sb.Append(Encoding.UTF8.GetString(buffer.Array!, 0, result.Count));
208 } while (!result.EndOfMessage);
209 return sb.ToString();
210 }
211}