From 94bb1bb93d23d83483a7ce3f7baf9caa5d4c3fb0 Mon Sep 17 00:00:00 2001
From: Abe <abe.malla8@gmail.com>
Date: Sun, 28 Jan 2024 14:07:54 -0800
Subject: [PATCH] Expose ServerInfo (#17)

* Expose ServerInfo

* Use InitializationResponse

* Remove spaces
---
 .../LanguageClient/InitializingServer.swift   | 74 ++++++++++++-------
 Sources/LanguageClient/RestartingServer.swift |  2 +-
 2 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/Sources/LanguageClient/InitializingServer.swift b/Sources/LanguageClient/InitializingServer.swift
index 4c736ab..fe72bc4 100644
--- a/Sources/LanguageClient/InitializingServer.swift
+++ b/Sources/LanguageClient/InitializingServer.swift
@@ -22,14 +22,14 @@ public actor InitializingServer {
 
 	private enum State {
 		case uninitialized
-		case initialized(ServerCapabilities)
+		case initialized(InitializationResponse)
 		case shutdown
 
 		var capabilities: ServerCapabilities? {
 			get {
 				switch self {
-				case .initialized(let capabilities):
-					return capabilities
+				case .initialized(let initResp):
+					return initResp.capabilities
 				case .uninitialized, .shutdown:
 					return nil
 				}
@@ -40,13 +40,22 @@ public actor InitializingServer {
 				}
 
 				switch self {
-				case .initialized:
-					self = .initialized(caps)
+				case .initialized(let initResp):
+					self = .initialized(initResp)
 				case .uninitialized, .shutdown:
 					break
 				}
 			}
 		}
+
+		var serverInfo: ServerInfo? {
+			switch self {
+			case .initialized(let initResp):
+				return initResp.serverInfo
+			case .uninitialized, .shutdown:
+				return nil
+			}
+		}
 	}
 
 	private let channel: ServerConnection
@@ -94,6 +103,23 @@ public actor InitializingServer {
 			return state.capabilities
 		}
 	}
+		
+	/// Return the server's info.
+	///
+	/// This will not start the server if it isn't already running.
+	public var serverInfo: ServerInfo? {
+		get async {
+			do {
+				try await semaphore.waitUnlessCancelled()
+			} catch {
+				return nil
+			}
+
+			defer { semaphore.signal() }
+
+			return state.serverInfo
+		}
+	}
 }
 
 extension InitializingServer: StatefulServer {
@@ -158,10 +184,10 @@ extension InitializingServer: StatefulServer {
 
 extension InitializingServer {
 	/// Run the initialization sequence with the server, if it has not already happened.
-	public func initializeIfNeeded() async throws -> ServerCapabilities {
+	public func initializeIfNeeded() async throws -> InitializationResponse {
 		switch state {
-		case .initialized(let caps):
-			return caps
+		case .initialized(let initResp):
+			return initResp
 		case .uninitialized, .shutdown:
 			try await semaphore.waitUnlessCancelled()
 		}
@@ -169,34 +195,31 @@ extension InitializingServer {
 		defer { semaphore.signal() }
 
 		let params = try await initializeParamsProvider()
-
 		let initResponse = try await channel.initialize(params)
-		let caps = initResponse.capabilities
 
 		try await channel.initialized(InitializedParams())
+		self.state = .initialized(initResponse)
 
-		self.state = .initialized(caps)
+		capabilitiesContinuation.yield(initResponse.capabilities)
 
-		capabilitiesContinuation.yield(caps)
-
-		return caps
+		return initResponse
 	}
 
 	private func handleEvent(_ event: ServerEvent) {
-    switch event {
-      case let .request(_, request):
-        handleRequest(request)
-      default:
-        break
-    }
-  }
+		switch event {
+		case let .request(_, request):
+			handleRequest(request)
+		default:
+			break
+		}
+	}
 
 	private func handleRequest(_ request: ServerRequest) {
-		guard case .initialized(let caps) = self.state else {
+		guard case .initialized(let initResp) = self.state else {
 			fatalError("received a request without being initialized")
 		}
 
-		var newCaps = caps
+		var newCaps = initResp.capabilities
 
 		do {
 			switch request {
@@ -211,9 +234,8 @@ extension InitializingServer {
 			print("unable to mutate server capabilities: \(error)")
 		}
 
-		if caps != newCaps {
-			self.state = .initialized(newCaps)
-
+		if initResp.capabilities != newCaps {
+			self.state = .initialized(InitializationResponse(capabilities: newCaps, serverInfo: initResp.serverInfo))
 			capabilitiesContinuation.yield(newCaps)
 		}
 	}
diff --git a/Sources/LanguageClient/RestartingServer.swift b/Sources/LanguageClient/RestartingServer.swift
index 0eaf051..b7bc74a 100644
--- a/Sources/LanguageClient/RestartingServer.swift
+++ b/Sources/LanguageClient/RestartingServer.swift
@@ -109,7 +109,7 @@ public actor RestartingServer<WrappedServer: ServerConnection & Sendable> {
 	}
 
 	/// Run the initialization sequence with the server, if it has not already happened.
-	public func initializeIfNeeded() async throws -> ServerCapabilities {
+	public func initializeIfNeeded() async throws -> InitializationResponse {
 		try await startServerIfNeeded().initializeIfNeeded()
 	}