diff --git a/ts/packages/agents/visualStudio/README.md b/ts/packages/agents/visualStudio/README.md index 4c430ec69a..e3e71eedcb 100644 --- a/ts/packages/agents/visualStudio/README.md +++ b/ts/packages/agents/visualStudio/README.md @@ -17,28 +17,51 @@ for build and install instructions. +-------------------------------+ +--------------------+ | Visual Studio (VSIX) | | TypeAgent | | | | agent-server | -| +------------------------+ | WS | | -| | ChatToolWindow |<--|------->| +--------------+ | +| +------------------------+ | | | +| | ChatToolWindow |<--|--WS--->| +--------------+ | | | WebView2 | | 8999 | | dispatcher | | | | chat-ui | | | +------+-------+ | | +------------------------+ | | | | | | | +------v-------+ | | +------------------------+ | | | visualstudio | | | | AgentBridgeClient.cs |<--|---WS---|->| agent | | -| | DTEActionExecutor.cs | | 5680 | +--------------+ | +| | DTEActionExecutor.cs | | ephem. | +--------------+ | | +-----------+------------+ | +--------------------+ -| | | -| v EnvDTE | -| (Solution, Build, | -| Debugger, Editor) | +| | ^ +| | + port discovered via +| | discovery channel on 8999 +| v EnvDTE +| (Solution, Build, +| Debugger, Editor) +-------------------------------+ ``` Two WebSocket channels: -- **Chat channel** (port 8999) — WebView2 inside the VSIX talks to the dispatcher. -- **Action bridge** (port 5680) — this agent owns a `WebSocketServer`; the C# - host connects as a client and dispatches incoming actions through EnvDTE. +- **Chat channel** (port 8999) — WebView2 inside the VSIX talks to the + dispatcher. The same port also hosts the dispatcher's read-only + **discovery channel** that the bridge uses to find the action port + (see below). +- **Action bridge** (OS-assigned ephemeral port) — this agent owns a + `WebSocketServer`; the C# host connects as a client and dispatches + incoming actions through EnvDTE. The actual port is published to the + agent-server's `PortRegistrar` under `(visualStudio, default)` and + discovered by `AgentBridgeClient` on every connect attempt. + +### Port discovery + +`AgentBridgeClient` resolves the bridge port on each connect attempt by +calling `lookupPort("visualStudio", "default")` against the +agent-server's discovery channel (`ws://localhost:/`, +default 8999). If discovery is unreachable or the agent isn't yet +registered, the reconnect loop simply retries — there is no silent +fallback to a well-known port. Knobs: + +| Env var | Default | Purpose | +| ----------------------------------------- | ------------------------- | ------------------------------------------------------------------------------------------------ | +| `AGENT_SERVER_PORT` | `8999` | Where the dispatcher's discovery channel is hosted. | +| `TYPEAGENT_VS_BRIDGE_PORT` _(C# host)_ | _(unset → use discovery)_ | Pin the C# host to dial a specific bridge port; bypasses discovery. Useful for manual debugging. | +| `VISUALSTUDIO_BRIDGE_PORT` _(agent-side)_ | _(unset → OS-assigned)_ | Pin the agent's bridge to a specific port when debugging. | ## Action Categories @@ -77,7 +100,8 @@ The VSIX host has its own two-stage build (WebView2 bundle + MSBuild). See 4. **View → Other Windows → TypeAgent Chat**. The chat panel connects to `ws://localhost:8999`. Once connected, the C# bridge -auto-connects to `ws://localhost:5680`. +discovers the agent's action port via that same WS (via the discovery channel) +and auto-connects. ## Project Structure @@ -126,8 +150,12 @@ Any use of third-party trademarks or logos are subject to those third-party's po - **Chat panel never connects.** The agent-server must be running before VS opens the tool window. Check that port 8999 is reachable. - **Actions hang or error with "Host plugin not connected".** The C# bridge - could not reach port 5680 — confirm the `visualstudio-agent` is enabled in - the dispatcher and that no other process is bound to 5680. + could not reach the agent's action port. Confirm `visualstudio-agent` is + enabled in the dispatcher. Check the VS **Output → Debug** pane for + `[TypeAgent] Bridge…` lines indicating whether discovery succeeded and + which port was used. To bypass discovery for diagnostics (e.g. dial a + manually-launched bridge), set `TYPEAGENT_VS_BRIDGE_PORT` to the port + you want to dial. - **Action ran but did nothing visible.** Check the Visual Studio **Output** window and the agent-server logs; EnvDTE silently no-ops on some commands when the relevant context (e.g. an active document, an active debug session) diff --git a/ts/packages/agents/visualStudio/host/csharp/Bridge/AgentBridgeClient.cs b/ts/packages/agents/visualStudio/host/csharp/Bridge/AgentBridgeClient.cs index bf1240bb9f..aaff568164 100644 --- a/ts/packages/agents/visualStudio/host/csharp/Bridge/AgentBridgeClient.cs +++ b/ts/packages/agents/visualStudio/host/csharp/Bridge/AgentBridgeClient.cs @@ -14,10 +14,21 @@ namespace Microsoft.TypeAgent.VisualStudio.Bridge { /// - /// WebSocket client that connects to the visualstudio-agent's bridge - /// (default ws://localhost:5680). Receives BridgeRequest messages, - /// dispatches them through DTEActionExecutor, and sends BridgeResponse - /// messages back. + /// WebSocket client that connects to the visualstudio-agent's bridge. + /// Receives BridgeRequest messages, dispatches them through + /// DTEActionExecutor, and sends BridgeResponse messages back. + /// + /// Port discovery: + /// The bridge port is no longer hardcoded. On each connect attempt + /// we ask the agent-server's discovery channel where the + /// `(visualStudio, default)` allocation lives. If discovery is + /// unreachable or the agent isn't yet registered, the reconnect + /// loop simply retries — there is no silent fallback to a + /// well-known port. To pin a specific port (e.g. when running the + /// bridge against a manually-launched agent), set + /// `TYPEAGENT_VS_BRIDGE_PORT`; that bypasses discovery entirely. + /// See for the wire protocol and the + /// `AGENT_SERVER_PORT` env-var knob. /// /// Wire format (matches packages/agents/visualStudio/src/visualStudioActionHandler.ts): /// request: { id, actionName, parameters } @@ -25,9 +36,7 @@ namespace Microsoft.TypeAgent.VisualStudio.Bridge /// internal sealed class AgentBridgeClient : IDisposable { - // Port 5678 + 5679 are taken by the Excel agent. Keep this in sync - // with BRIDGE_PORT in packages/agents/visualStudio/src/visualStudioActionHandler.ts. - private static readonly Uri DefaultUri = new Uri("ws://localhost:5680"); + private const string BridgePortOverrideEnv = "TYPEAGENT_VS_BRIDGE_PORT"; private static readonly TimeSpan ReconnectDelay = TimeSpan.FromSeconds(3); private readonly AsyncPackage _package; @@ -46,9 +55,26 @@ public async Task StartAsync(CancellationToken cancellation) using var linked = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, cancellation); while (!linked.IsCancellationRequested) { + int port = 0; try { - await ConnectAndReceiveAsync(linked.Token); + // Resolve the port fresh on every attempt: the agent may + // have restarted on a different ephemeral port since the + // last loop iteration, and the standalone shell may have + // come up while we were retrying. + int? resolved = ResolvePortOverride() + ?? await BridgeDiscovery.ResolveBridgePortAsync(linked.Token).ConfigureAwait(false); + if (resolved is null) + { + // Discovery succeeded but the agent isn't registered + // yet — wait one reconnect cycle and try again. + Debug.WriteLine("[TypeAgent] visualStudio agent not yet registered; will retry"); + } + else + { + port = resolved.Value; + await ConnectAndReceiveAsync(port, linked.Token).ConfigureAwait(false); + } } catch (OperationCanceledException) { @@ -56,11 +82,11 @@ public async Task StartAsync(CancellationToken cancellation) } catch (Exception ex) { - Debug.WriteLine($"[TypeAgent] Bridge error: {ex.Message}"); + Debug.WriteLine($"[TypeAgent] Bridge error (port {port}): {ex.Message}"); } try { - await Task.Delay(ReconnectDelay, linked.Token); + await Task.Delay(ReconnectDelay, linked.Token).ConfigureAwait(false); } catch (OperationCanceledException) { @@ -69,11 +95,28 @@ public async Task StartAsync(CancellationToken cancellation) } } - private async Task ConnectAndReceiveAsync(CancellationToken cancellation) + // Returns an explicit port override from `TYPEAGENT_VS_BRIDGE_PORT`, + // or null when the env var is unset/malformed (caller falls through + // to discovery). Mirrors `CODE_WEBSOCKET_HOST` from coda. + private static int? ResolvePortOverride() + { + string? raw = Environment.GetEnvironmentVariable(BridgePortOverrideEnv); + if (string.IsNullOrEmpty(raw)) return null; + if (int.TryParse(raw, out int p) && p > 0 && p <= 65535) + { + Debug.WriteLine($"[TypeAgent] {BridgePortOverrideEnv} override active: {p}"); + return p; + } + Debug.WriteLine($"[TypeAgent] Ignoring malformed {BridgePortOverrideEnv}={raw}"); + return null; + } + + private async Task ConnectAndReceiveAsync(int port, CancellationToken cancellation) { + var uri = new Uri($"ws://localhost:{port}"); _ws = new ClientWebSocket(); - await _ws.ConnectAsync(DefaultUri, cancellation); - Debug.WriteLine($"[TypeAgent] Bridge connected to {DefaultUri}"); + await _ws.ConnectAsync(uri, cancellation).ConfigureAwait(false); + Debug.WriteLine($"[TypeAgent] Bridge connected to {uri}"); var buffer = new ArraySegment(new byte[16 * 1024]); var assembly = new StringBuilder(); @@ -84,10 +127,10 @@ private async Task ConnectAndReceiveAsync(CancellationToken cancellation) WebSocketReceiveResult result; do { - result = await _ws.ReceiveAsync(buffer, cancellation); + result = await _ws.ReceiveAsync(buffer, cancellation).ConfigureAwait(false); if (result.MessageType == WebSocketMessageType.Close) { - await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cancellation); + await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cancellation).ConfigureAwait(false); return; } assembly.Append(Encoding.UTF8.GetString(buffer.Array!, 0, result.Count)); diff --git a/ts/packages/agents/visualStudio/host/csharp/Bridge/BridgeDiscovery.cs b/ts/packages/agents/visualStudio/host/csharp/Bridge/BridgeDiscovery.cs new file mode 100644 index 0000000000..63c34069f9 --- /dev/null +++ b/ts/packages/agents/visualStudio/host/csharp/Bridge/BridgeDiscovery.cs @@ -0,0 +1,170 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System; +using System.Diagnostics; +using System.Net.WebSockets; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; + +namespace Microsoft.TypeAgent.VisualStudio.Bridge +{ + /// + /// Looks up the visualStudio agent's bridge port via the dispatcher's + /// discovery channel. + /// + /// Wire protocol (matches packages/agentServer/protocol/src/protocol.ts): + /// client → server: + /// { "name": "discovery", + /// "message": { "type": "invoke", "callId": N, "name": "lookupPort", + /// "args": [{ "agentName": "visualStudio", + /// "role": "default" }] } } + /// server → client: + /// { "name": "discovery", + /// "message": { "type": "invokeResult", "callId": N, + /// "result": { "port": <int|null> } } } + /// + /// Returns the discovered port, or null when the agent isn't yet + /// registered with the agent-server (transient — caller should retry). + /// Throws on transport failure so the outer reconnect loop can apply + /// its own retry/backoff. There is intentionally no hardcoded fallback + /// port — the migrated TS clients (browser, code, coda) all return + /// "undefined" on discovery failure and rely on the reconnect loop; + /// dialing a stale well-known port would just connect to nothing. + /// + internal static class BridgeDiscovery + { + // Read on every resolve so users can flip behavior without + // restarting the IDE between debugging sessions. + private const string AgentServerPortEnv = "AGENT_SERVER_PORT"; + + // Must match AGENT_SERVER_DEFAULT_PORT in agentServer/protocol. + private const int DefaultAgentServerPort = 8999; + + // Names this client uses to look itself up. Must match the role + // registered by visualStudioActionHandler.ts. + private const string AgentName = "visualStudio"; + private const string Role = "default"; + + /// + /// Resolve the bridge port via discovery. Returns the discovered + /// port, or null when the agent has not yet registered (transient + /// — caller should retry on its reconnect loop). + /// Throws on transport failure (agent-server unreachable, timeout, + /// malformed response) so the caller can log and retry. + /// + public static async Task ResolveBridgePortAsync(CancellationToken cancellation) + { + int agentServerPort = GetAgentServerPort(); + int? discovered = await LookupPortAsync(agentServerPort, cancellation).ConfigureAwait(false); + if (discovered is int p) + { + Debug.WriteLine($"[TypeAgent] Discovery resolved bridge port {p}"); + } + else + { + Debug.WriteLine($"[TypeAgent] Discovery returned null for ({AgentName}, {Role}); agent not yet registered"); + } + return discovered; + } + + private static int GetAgentServerPort() + { + string? raw = Environment.GetEnvironmentVariable(AgentServerPortEnv); + if (int.TryParse(raw, out int p) && p > 0 && p <= 65535) + { + return p; + } + return DefaultAgentServerPort; + } + + private static async Task LookupPortAsync(int agentServerPort, CancellationToken cancellation) + { + var uri = new Uri($"ws://localhost:{agentServerPort}/"); + using var ws = new ClientWebSocket(); + // Cap the discovery call so a hung agent-server doesn't stall + // the whole reconnect loop. The outer AgentBridgeClient loop + // already retries on a separate cadence. + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5)); + using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellation, timeout.Token); + + await ws.ConnectAsync(uri, linked.Token).ConfigureAwait(false); + + // callId is arbitrary — the server echoes it back verbatim, + // and we only have one outstanding request per socket. + const int callId = 1; + var request = new JObject + { + ["name"] = "discovery", + ["message"] = new JObject + { + ["type"] = "invoke", + ["callId"] = callId, + ["name"] = "lookupPort", + ["args"] = new JArray + { + new JObject + { + ["agentName"] = AgentName, + ["role"] = Role, + }, + }, + }, + }; + byte[] requestBytes = Encoding.UTF8.GetBytes(request.ToString(Formatting.None)); + await ws.SendAsync( + new ArraySegment(requestBytes), + WebSocketMessageType.Text, + endOfMessage: true, + linked.Token).ConfigureAwait(false); + + string responseText = await ReceiveFullMessageAsync(ws, linked.Token).ConfigureAwait(false); + try + { + await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, CancellationToken.None).ConfigureAwait(false); + } + catch + { + // Best-effort close — the response is already in hand. + } + + var root = JObject.Parse(responseText); + string? name = root.Value("name"); + if (name != "discovery") return null; + var inner = root["message"] as JObject; + if (inner == null) return null; + string? type = inner.Value("type"); + if (type == "invokeError") + { + throw new InvalidOperationException( + inner.Value("error") ?? "Discovery returned invokeError"); + } + if (type != "invokeResult") return null; + if (inner.Value("callId") != callId) return null; + var result = inner["result"] as JObject; + if (result == null) return null; + // `port` is `int|null`; JObject returns null cleanly for both. + return result.Value("port"); + } + + private static async Task ReceiveFullMessageAsync(ClientWebSocket ws, CancellationToken cancellation) + { + var buffer = new ArraySegment(new byte[16 * 1024]); + var sb = new StringBuilder(); + WebSocketReceiveResult result; + do + { + result = await ws.ReceiveAsync(buffer, cancellation).ConfigureAwait(false); + if (result.MessageType == WebSocketMessageType.Close) + { + throw new InvalidOperationException("Discovery WS closed before response"); + } + sb.Append(Encoding.UTF8.GetString(buffer.Array!, 0, result.Count)); + } while (!result.EndOfMessage); + return sb.ToString(); + } + } +} diff --git a/ts/packages/agents/visualStudio/host/csharp/VisualStudioTypeAgent.csproj b/ts/packages/agents/visualStudio/host/csharp/VisualStudioTypeAgent.csproj index a179d13c6d..866c8f99e7 100644 --- a/ts/packages/agents/visualStudio/host/csharp/VisualStudioTypeAgent.csproj +++ b/ts/packages/agents/visualStudio/host/csharp/VisualStudioTypeAgent.csproj @@ -73,6 +73,7 @@ ChatToolWindowControl.xaml + diff --git a/ts/packages/agents/visualStudio/package.json b/ts/packages/agents/visualStudio/package.json index 166a51b414..2535cf1281 100644 --- a/ts/packages/agents/visualStudio/package.json +++ b/ts/packages/agents/visualStudio/package.json @@ -25,10 +25,13 @@ }, "dependencies": { "@typeagent/agent-sdk": "workspace:*", + "debug": "^4.4.0", + "websocket-utils": "workspace:*", "ws": "^8.20.1" }, "devDependencies": { "@typeagent/action-schema-compiler": "workspace:*", + "@types/debug": "^4.1.12", "@types/ws": "^8.5.10", "action-grammar-compiler": "workspace:*", "concurrently": "^9.1.2", diff --git a/ts/packages/agents/visualStudio/src/originAllowlist.ts b/ts/packages/agents/visualStudio/src/originAllowlist.ts new file mode 100644 index 0000000000..737c289385 --- /dev/null +++ b/ts/packages/agents/visualStudio/src/originAllowlist.ts @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { createAgentOriginAllowlist } from "websocket-utils/originAllowlist"; + +/** + * Origin allowlist for the visualStudio agent's WebSocket bridge. + * + * The only legitimate client is the in-process Visual Studio extension + * (`host/csharp/Bridge/AgentBridgeClient.cs`), which uses + * `System.Net.WebSockets.ClientWebSocket`. That client does **not** + * send an `Origin` header, so we rely on the shared no-Origin baseline + * documented on {@link createAgentOriginAllowlist}. + * + * No extension scheme prefixes are accepted — anything beyond loopback + * web clients (manual debugging, future webview consumers) and + * Origin-less native clients is rejected with HTTP 403 before the + * `connection` event fires. Every per-agent listener that binds to an + * ephemeral port via the PortRegistrar must gate Origin so those ports + * can't be dialed by arbitrary web pages on the same host. + */ +export const isAllowedAgentOrigin = createAgentOriginAllowlist(); diff --git a/ts/packages/agents/visualStudio/src/visualStudioActionHandler.ts b/ts/packages/agents/visualStudio/src/visualStudioActionHandler.ts index 6f85ec3ecb..762293a026 100644 --- a/ts/packages/agents/visualStudio/src/visualStudioActionHandler.ts +++ b/ts/packages/agents/visualStudio/src/visualStudioActionHandler.ts @@ -4,6 +4,15 @@ // Pattern: websocket-bridge — bidirectional RPC to a host-side plugin. // The agent owns a WebSocketServer; the host plugin connects as the client. // Commands flow TypeAgent → WebSocket → plugin → response. +// +// Port handling: the bridge binds an ephemeral port (OS-assigned via +// `bind(0)`) and publishes the actual port to the PortRegistrar under +// role `"default"`. The Visual Studio extension discovers the port via +// the dispatcher's discovery channel (`AGENT_SERVER_PORT` → +// `lookupPort("visualStudio", "default")`). For local debugging the env +// var `VISUALSTUDIO_BRIDGE_PORT` pins a fixed port (useful when running +// the extension in a debugger that can't easily read the discovery +// channel). import { ActionContext, @@ -13,12 +22,13 @@ import { ActionResult, } from "@typeagent/agent-sdk"; import { createActionResultFromTextDisplay } from "@typeagent/agent-sdk/helpers/action"; -import { WebSocketServer, WebSocket } from "ws"; +import { WebSocketServer, WebSocket, RawData } from "ws"; +import { AddressInfo } from "net"; +import registerDebug from "debug"; import { VisualStudioActions } from "./visualStudioSchema.js"; +import { isAllowedAgentOrigin } from "./originAllowlist.js"; -// Port 5678 + 5679 are taken by the Excel agent (in the SecretAgents repo). -// Keep this in sync with AgentBridgeClient.cs in the VS extension host. -const BRIDGE_PORT = 5680; +const debug = registerDebug("typeagent:visualstudio:bridge"); // ---- WebSocket bridge -------------------------------------------------- @@ -31,84 +41,428 @@ type BridgeResponse = { }; class VisualStudioBridge { - private wss: WebSocketServer | undefined; - private client: WebSocket | undefined; - private pending = new Map void>(); - - start(): void { - const wss = new WebSocketServer({ port: BRIDGE_PORT }); - wss.on("error", (err) => { - console.error( - `[visualStudio] bridge WebSocketServer failed on port ${BRIDGE_PORT}:`, - err, - ); + // Multi-client tracking. The bridge accepts every VS extension that + // dials in (insertion-order-preserving Map so we can route to the + // most-recently-connected one). Keys are synthetic ids — the VS + // extension's `ClientWebSocket` doesn't expose anything we could + // use to distinguish two parallel instances. + private clients: Map = new Map(); + private clientIdCounter = 0; + // For each in-flight request we record which clientId we sent it + // to, so a disconnect on one VS only fails the requests that + // actually targeted that instance — other VS instances keep + // running. The id itself is bridge-wide (responses are looked up + // by id alone). + private pending = new Map< + string, + { + resolve: (r: BridgeResponse) => void; + timer: NodeJS.Timeout; + clientId: string; + } + >(); + private readonly sendTimeoutMs: number; + + /** + * Invoked on every client connect/disconnect with the current count + * of OPEN clients. Set by the agent-lifecycle code; the bridge + * itself doesn't care about consumers. Errors raised by the + * callback are swallowed (the count is informational and only + * surfaces via `@system ports`). + */ + public onClientCountChanged?: (count: number) => void; + + /** + * @param wss the underlying ws server, already bound and listening. + * @param port the actually bound port (OS-assigned when caller + * passed 0, or the env-pinned value when set). + * + * Construction is private — use {@link VisualStudioBridge.start} so + * callers always get a bridge that is guaranteed to be bound before + * they read {@link port} or pass it to the registrar. + */ + private constructor( + private readonly wss: WebSocketServer, + public readonly port: number, + ) { + this.sendTimeoutMs = resolveSendTimeoutMs(); + this.setupHandlers(); + debug( + `VisualStudioBridge listening on port ${port} (sendTimeoutMs=${this.sendTimeoutMs})`, + ); + } + + /** + * Bind a new bridge on `port`. Resolves only after the `listening` + * event so callers can synchronously read {@link port}; rejects on + * the first `error` event so bind failures (EADDRINUSE under a + * fixed-port override) surface loudly instead of being swallowed by + * an attached error handler. + * + * Pass `0` to let the OS pick a free ephemeral port; the actual port + * is then available via {@link port}. + */ + public static start(port: number = 0): Promise { + return new Promise((resolve, reject) => { + const wss = new WebSocketServer({ + port, + // Bind loopback-only so the bridge isn't reachable from + // other hosts on the LAN. The Origin allowlist below + // accepts requests with no Origin header (the C# + // ClientWebSocket doesn't send one), so without an + // explicit loopback bind a remote attacker on the same + // network could otherwise drive EnvDTE actions. + host: "127.0.0.1", + // Gate every upgrade on Origin so a random web page on + // the same host can't dial the ephemeral port assigned + // by the OS. `verifyClient` runs synchronously before + // the `connection` event fires; rejected requests get + // HTTP 403. The C# `ClientWebSocket` doesn't send + // Origin, which `isAllowedAgentOrigin` accepts. + verifyClient: (info, cb) => { + const origin = info.req.headers.origin as + | string + | undefined; + if (isAllowedAgentOrigin(origin)) { + cb(true); + } else { + debug(`Rejecting WS upgrade from origin ${origin}`); + cb(false, 403, "Origin not allowed"); + } + }, + }); + let settled = false; + const onError = (error: Error) => { + if (settled) { + debug("Server error after listening:", error); + return; + } + settled = true; + wss.removeListener("listening", onListening); + debug("Server bind error:", error); + reject(error); + }; + const onListening = () => { + if (settled) return; + settled = true; + wss.removeListener("error", onError); + const address = wss.address() as AddressInfo | null; + if (!address || typeof address === "string") { + wss.close(); + reject( + new Error( + "ws server.address() did not return an AddressInfo", + ), + ); + return; + } + // Re-attach a permanent error handler so post-listen + // errors are logged rather than crashing the process. + wss.on("error", (err) => { + debug("Server error:", err); + }); + resolve(new VisualStudioBridge(wss, address.port)); + }; + wss.once("error", onError); + wss.once("listening", onListening); }); - wss.on("listening", () => { - console.log( - `[visualStudio] bridge listening on ws://localhost:${BRIDGE_PORT}`, + } + + private setupHandlers(): void { + this.wss.on("connection", (ws: WebSocket) => { + const clientId = `vs-${++this.clientIdCounter}`; + this.clients.set(clientId, ws); + debug( + `host plugin connected (${clientId}); total=${this.clients.size}`, ); - }); - wss.on("connection", (ws) => { - console.log("[visualStudio] host plugin connected"); - this.client = ws; - ws.on("message", (data) => { - const response = JSON.parse(data.toString()) as BridgeResponse; - this.pending.get(response.id)?.(response); - this.pending.delete(response.id); - }); - ws.on("close", () => { - console.log("[visualStudio] host plugin disconnected"); - this.client = undefined; + this.emitClientCount(); + ws.on("message", (data: RawData) => { + try { + const response = JSON.parse( + data.toString(), + ) as BridgeResponse; + const entry = this.pending.get(response.id); + if (entry !== undefined) { + clearTimeout(entry.timer); + this.pending.delete(response.id); + entry.resolve(response); + } + } catch (err) { + debug("Failed to parse plugin message:", err); + } }); + const onDisconnect = (reason: string) => { + // Identity-guarded so the second of {close, error} is a + // no-op (Map.delete returns false the second time). + if (!this.clients.delete(clientId)) return; + debug( + `host plugin disconnected (${clientId}, ${reason}); remaining=${this.clients.size}`, + ); + this.emitClientCount(); + // Only fail the pending sends that targeted *this* + // client — other VS instances stay live and their + // pending requests must keep waiting for their own + // responses (or hit the per-request timeout). + this.failPendingForClient( + clientId, + new Error(`Host plugin disconnected: ${reason}`), + ); + }; + ws.on("close", () => onDisconnect("close")); ws.on("error", (err) => { - console.error("[visualStudio] host plugin socket error:", err); + debug("host plugin socket error:", err); + onDisconnect(`error: ${err.message}`); }); }); - this.wss = wss; } + private emitClientCount(): void { + try { + this.onClientCountChanged?.(this.getConnectedCount()); + } catch (err) { + debug("onClientCountChanged threw:", err); + } + } + + /** + * Number of currently-OPEN bridge clients. Surfaced via + * `@system ports` through the SessionContext's + * `notifyClientCountChanged` API. Multiple VS instances can + * connect concurrently and each is counted independently. + */ + public getConnectedCount(): number { + let n = 0; + for (const ws of this.clients.values()) { + if (ws.readyState === WebSocket.OPEN) n++; + } + return n; + } + + private failPending(error: Error): void { + if (this.pending.size === 0) return; + const entries = Array.from(this.pending.values()); + this.pending.clear(); + for (const entry of entries) { + clearTimeout(entry.timer); + entry.resolve({ + id: "", + success: false, + error: error.message, + }); + } + } + + private failPendingForClient(clientId: string, error: Error): void { + const toFail: string[] = []; + for (const [id, entry] of this.pending) { + if (entry.clientId === clientId) toFail.push(id); + } + for (const id of toFail) { + const entry = this.pending.get(id)!; + this.pending.delete(id); + clearTimeout(entry.timer); + entry.resolve({ + id: "", + success: false, + error: error.message, + }); + } + } + + /** + * Close the underlying server and resolve once the port is fully + * released — important for a rapid disable→enable cycle under a + * fixed-port override (`VISUALSTUDIO_BRIDGE_PORT`), where a + * synchronous return would race the new bind into EADDRINUSE. + */ async stop(): Promise { - return new Promise((resolve) => this.wss?.close(() => resolve())); + debug("Closing VisualStudioBridge"); + for (const ws of this.clients.values()) { + if (ws.readyState === WebSocket.OPEN) { + ws.close(); + } + } + this.clients.clear(); + // Reject any in-flight sends before the server closes; otherwise + // callers awaiting `send()` hang forever after a manual disable. + this.failPending(new Error("VisualStudioBridge stopped")); + return new Promise((resolve) => this.wss.close(() => resolve())); } async send(actionName: string, parameters: unknown): Promise { - if (!this.client) { - throw new Error("No host plugin connected on port " + BRIDGE_PORT); + // Route to the most-recently-connected OPEN client. This + // preserves the legacy single-client behavior (where the + // newest connection won) when only one VS is active, and + // gives a sensible "last-wins" tiebreak for multi-VS sessions + // — RPCs like buildSolution can't sensibly broadcast (both + // VS instances would execute), and a smarter routing strategy + // (per-solution-path, user-disambiguated) is future work. + const target = this.pickTargetClient(); + if (target === undefined) { + throw new Error(`No host plugin connected on port ${this.port}`); } + const [clientId, ws] = target; const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`; return new Promise((resolve, reject) => { - this.pending.set(id, (res) => - res.success - ? resolve(res.result) - : reject(new Error(res.error)), - ); - this.client!.send( - JSON.stringify({ - id, - actionName, - parameters, - } satisfies BridgeRequest), - ); + // Timeout so a plugin that accepts the WS frame but never + // replies (deadlock, EnvDTE hang, killed mid-action) doesn't + // wedge the caller. The pending entry is cleared either by + // the response handler or this timer — whichever fires + // first. + const timer = setTimeout(() => { + if (this.pending.delete(id)) { + reject( + new Error( + `VS bridge action '${actionName}' timed out after ${this.sendTimeoutMs}ms`, + ), + ); + } + }, this.sendTimeoutMs); + this.pending.set(id, { + timer, + clientId, + resolve: (res) => + res.success + ? resolve(res.result) + : reject(new Error(res.error)), + }); + try { + ws.send( + JSON.stringify({ + id, + actionName, + parameters, + } satisfies BridgeRequest), + ); + } catch (err) { + clearTimeout(timer); + this.pending.delete(id); + reject(err); + } }); } + // Map iteration is insertion-ordered; walk in reverse so the + // newest OPEN connection wins. Returns undefined when no client + // is currently OPEN. + private pickTargetClient(): [string, WebSocket] | undefined { + const entries = Array.from(this.clients.entries()); + for (let i = entries.length - 1; i >= 0; i--) { + const [id, ws] = entries[i]; + if (ws.readyState === WebSocket.OPEN) return [id, ws]; + } + return undefined; + } + get connected(): boolean { - return this.client !== undefined; + return this.getConnectedCount() > 0; } } +// ---- Port resolution --------------------------------------------------- + +const DEFAULT_SEND_TIMEOUT_MS = 30_000; + +// Per-action send() timeout. EnvDTE actions are typically subsecond, +// but a few (build/run, attach-to-process) can run for tens of seconds +// on a cold solution. Default 30s leaves headroom; override via env for +// debugging long-running actions. +function resolveSendTimeoutMs(): number { + const raw = process.env["VISUALSTUDIO_BRIDGE_SEND_TIMEOUT_MS"]; + if (raw === undefined || raw === "") return DEFAULT_SEND_TIMEOUT_MS; + const n = Number(raw); + if (!Number.isInteger(n) || n <= 0) { + console.warn( + `Ignoring malformed VISUALSTUDIO_BRIDGE_SEND_TIMEOUT_MS=${raw}; using ${DEFAULT_SEND_TIMEOUT_MS}ms`, + ); + return DEFAULT_SEND_TIMEOUT_MS; + } + return n; +} + +// Optional fixed-port override. Useful when launching the Visual Studio +// extension in a debugger and you want both sides on a known port — +// pair this with `TYPEAGENT_VS_BRIDGE_PORT=` on the C# host +// to bypass discovery on both ends. +// +// Malformed values are warned and ignored — we fall through to the +// OS-assigned port. If the requested port is already in use, +// `VisualStudioBridge.start()` rejects with that error instead of +// silently rebinding. +function getBridgeBindPort(): number { + const raw = process.env["VISUALSTUDIO_BRIDGE_PORT"]; + if (raw === undefined || raw === "") return 0; + const n = Number(raw); + if (!Number.isInteger(n) || n < 0 || n > 65535) { + console.warn( + `Ignoring malformed VISUALSTUDIO_BRIDGE_PORT=${raw}; using OS-assigned port instead`, + ); + return 0; + } + debug(`VISUALSTUDIO_BRIDGE_PORT override active: ${n}`); + return n; +} + // ---- Agent lifecycle --------------------------------------------------- -type Context = { bridge: VisualStudioBridge }; +type Context = { + bridge?: VisualStudioBridge; + portRegistration?: { release: () => void }; +}; // Shared, process-singleton bridge: the VS extension only ever opens one -// WebSocket connection on port 5680, so per-session bridges would collide on -// `new WebSocketServer({ port: 5680 })` (EADDRINUSE) the moment a second -// session/conversation initializes the visualStudio agent. Created on first -// initialize, closed when the last session disables. Mirrors the pattern used -// by the code agent's CodeAgentWebSocketServer. +// WebSocket connection on the bound port, so per-session bridges would +// collide on `new WebSocketServer({ port })` (EADDRINUSE under fixed-port +// override) the moment a second session/conversation initializes the +// visualStudio agent. Created on first initialize, closed when the last +// session disables. Mirrors the pattern used by the code agent's +// `CodeAgentWebSocketServer`. let sharedBridge: VisualStudioBridge | undefined; +let sharedStartingPromise: Promise | undefined; +let sharedClosingPromise: Promise | undefined; let sharedBridgeRefCount = 0; +// Active sessions currently holding a `(visualStudio, default)` +// registration on the shared bridge. Insertion order picks a "primary" +// session for client-count reporting: the primary publishes the global +// count, the rest publish 0, so `@system ports` doesn't double-count +// when summing per-session entries. See codeActionHandler for the +// reference implementation. +const sharedActiveSessions = new Set>(); + +function publishClientCountFanout(count: number): void { + const primary = sharedActiveSessions.values().next().value; + for (const sc of sharedActiveSessions) { + void sc.notifyClientCountChanged("default", sc === primary ? count : 0); + } +} + +async function ensureSharedBridge(): Promise { + // If a previous teardown is still releasing the port, await it + // before binding again (matters under VISUALSTUDIO_BRIDGE_PORT + // override). + if (sharedClosingPromise !== undefined) { + await sharedClosingPromise; + } + if (sharedBridge !== undefined) return sharedBridge; + if (sharedStartingPromise !== undefined) return sharedStartingPromise; + sharedStartingPromise = (async () => { + try { + const bridge = await VisualStudioBridge.start(getBridgeBindPort()); + // Fan out client-count updates to active sessions. The + // bridge fires this on every connect/disconnect; the + // primary-session pattern in `publishClientCountFanout` + // prevents the `@system ports` summing logic from + // double-counting across sessions that each registered + // the SAME physical port. + bridge.onClientCountChanged = publishClientCountFanout; + sharedBridge = bridge; + return bridge; + } finally { + sharedStartingPromise = undefined; + } + })(); + return sharedStartingPromise; +} export function instantiate(): AppAgent { return { @@ -120,32 +474,96 @@ export function instantiate(): AppAgent { } async function initializeAgentContext(): Promise { - if (sharedBridge === undefined) { - sharedBridge = new VisualStudioBridge(); - sharedBridge.start(); - } - sharedBridgeRefCount++; - return { bridge: sharedBridge }; + return {}; } async function updateAgentContext( - _enable: boolean, - _context: SessionContext, + enable: boolean, + context: SessionContext, _schemaName: string, -): Promise {} +): Promise { + const agentContext = context.agentContext; + if (enable) { + if (agentContext.bridge !== undefined) return; + try { + const bridge = await ensureSharedBridge(); + agentContext.bridge = bridge; + // Per-session registration: the registrar allows multiple + // entries for `(visualStudio, default)` across sessions and + // lookup returns the most recent, so each active session + // independently keeps the shared port discoverable. The + // backstop in closeSessionContext releases ours if disable + // is skipped. + agentContext.portRegistration = context.registerPort( + "default", + bridge.port, + ); + sharedBridgeRefCount++; + sharedActiveSessions.add(context); + // Publish the current (global) count to the primary + // session (first in insertion order) and 0 to this session + // if it isn't the primary, so `@system ports` summing + // doesn't double-count. If this session is now becoming + // the primary (i.e. it's the first to enable), it gets the + // real count; otherwise it reports 0 and any future + // onClientCountChanged fanout keeps it at 0. + const primary = sharedActiveSessions.values().next().value; + void context.notifyClientCountChanged( + "default", + context === primary ? bridge.getConnectedCount() : 0, + ); + } catch (e) { + // Roll back per-session bookkeeping so a subsequent retry + // sees a clean slate. The shared bridge is left untouched — + // its refcount only advances on the success path below + // `registerPort`, and other sessions may still be using it. + delete agentContext.bridge; + delete agentContext.portRegistration; + throw e; + } + } else { + if (agentContext.bridge === undefined) return; + delete agentContext.bridge; + // Release this session's registration before potentially closing + // the bridge. Release is idempotent and a no-op if already + // released by the backstop. + agentContext.portRegistration?.release(); + delete agentContext.portRegistration; + + const wasPrimary = + sharedActiveSessions.values().next().value === context; + sharedActiveSessions.delete(context); + + sharedBridgeRefCount = Math.max(0, sharedBridgeRefCount - 1); + if (sharedBridgeRefCount === 0 && sharedBridge !== undefined) { + const toStop = sharedBridge; + sharedBridge = undefined; + // Track the in-flight close so a rapid re-enable awaits + // port release under a fixed-port override. + sharedClosingPromise = toStop.stop().finally(() => { + sharedClosingPromise = undefined; + }); + await sharedClosingPromise; + } else if (wasPrimary && sharedBridge !== undefined) { + // Primary session went away — transfer the (global) count + // to the new primary so `@system ports` keeps reporting + // the real number instead of 0. + const newPrimary = sharedActiveSessions.values().next().value; + if (newPrimary !== undefined) { + void newPrimary.notifyClientCountChanged( + "default", + sharedBridge.getConnectedCount(), + ); + } + } + } +} async function closeAgentContext( - _context: SessionContext, + context: SessionContext, ): Promise { - if (sharedBridgeRefCount === 0) { - return; - } - sharedBridgeRefCount--; - if (sharedBridgeRefCount === 0 && sharedBridge !== undefined) { - const toStop = sharedBridge; - sharedBridge = undefined; - await toStop.stop(); - } + // Defensive cleanup if updateAgentContext(false) wasn't invoked. + await updateAgentContext(false, context, ""); } async function executeAction( @@ -153,9 +571,14 @@ async function executeAction( context: ActionContext, ): Promise { const { bridge } = context.sessionContext.agentContext; + if (!bridge) { + return { + error: "visualStudio agent is not enabled in this session.", + }; + } if (!bridge.connected) { return { - error: `Host plugin not connected. Make sure the visualStudio plugin is running on port ${BRIDGE_PORT}.`, + error: `Host plugin not connected on port ${bridge.port}. Make sure the visualStudio extension is running.`, }; } try { diff --git a/ts/pnpm-lock.yaml b/ts/pnpm-lock.yaml index 0c356272b0..5789b454e6 100644 --- a/ts/pnpm-lock.yaml +++ b/ts/pnpm-lock.yaml @@ -3251,6 +3251,12 @@ importers: '@typeagent/agent-sdk': specifier: workspace:* version: link:../../agentSdk + debug: + specifier: ^4.4.0 + version: 4.4.3(supports-color@8.1.1) + websocket-utils: + specifier: workspace:* + version: link:../../utils/webSocketUtils ws: specifier: ^8.20.1 version: 8.20.1 @@ -3258,6 +3264,9 @@ importers: '@typeagent/action-schema-compiler': specifier: workspace:* version: link:../../actionSchemaCompiler + '@types/debug': + specifier: ^4.1.12 + version: 4.1.12 '@types/ws': specifier: ^8.5.10 version: 8.18.1 @@ -17694,7 +17703,7 @@ snapshots: '@babel/types': 7.28.4 '@jridgewell/remapping': 2.3.5 convert-source-map: 2.0.0 - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) gensync: 1.0.0-beta.2 json5: 2.2.3 semver: 6.3.1 @@ -22447,7 +22456,7 @@ snapshots: builder-util-runtime: 9.5.1 chromium-pickle-js: 0.2.0 ci-info: 4.3.1 - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) dmg-builder: 26.8.1(electron-builder-squirrel-windows@26.8.1) dotenv: 16.5.0 dotenv-expand: 11.0.7 @@ -22867,14 +22876,14 @@ snapshots: builder-util-runtime@9.3.1: dependencies: - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) sax: 1.3.0 transitivePeerDependencies: - supports-color builder-util-runtime@9.5.1: dependencies: - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) sax: 1.3.0 transitivePeerDependencies: - supports-color @@ -22887,7 +22896,7 @@ snapshots: builder-util-runtime: 9.5.1 chalk: 4.1.2 cross-spawn: 7.0.6 - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) fs-extra: 10.1.0 http-proxy-agent: 7.0.2 https-proxy-agent: 7.0.6 @@ -24120,7 +24129,7 @@ snapshots: electron-winstaller@5.4.0: dependencies: '@electron/asar': 3.4.1 - debug: 4.4.1 + debug: 4.4.3(supports-color@8.1.1) fs-extra: 7.0.1 lodash: 4.18.1 temp: 0.9.4