Fix that the request may timeout too soon

Author: intitni

Tool/Sources/GitHubCopilotService/LanguageServer/GitHubCopilotService.swift

@@ -37,21 +37,33 @@ public protocol GitHubCopilotSuggestionServiceType {
 }
 
 protocol GitHubCopilotLSP {
-    func sendRequest<E: GitHubCopilotRequestType>(_ endpoint: E) async throws -> E.Response
+    func sendRequest<E: GitHubCopilotRequestType>(
+        _ endpoint: E,
+        timeout: TimeInterval?
+    ) async throws -> E.Response
     func sendNotification(_ notif: ClientNotification) async throws
 }
 
+extension GitHubCopilotLSP {
+    func sendRequest<E: GitHubCopilotRequestType>(_ endpoint: E) async throws -> E.Response {
+        try await sendRequest(endpoint, timeout: nil)
+    }
+}
+
 enum GitHubCopilotError: Error, LocalizedError {
     case languageServerNotInstalled
     case languageServerError(ServerError)
     case failedToInstallStartScript
+    case chatEndsWithError(String)
 
     var errorDescription: String? {
         switch self {
         case .languageServerNotInstalled:
             return "Language server is not installed."
         case .failedToInstallStartScript:
             return "Failed to install start script."
+        case let .chatEndsWithError(errorMessage):
+            return "Chat ended with error message: \(errorMessage)"
         case let .languageServerError(error):
             switch error {
             case let .handlerUnavailable(handler):
@@ -578,8 +590,19 @@ public final class GitHubCopilotService: GitHubCopilotBaseService,
 }
 
 extension InitializingServer: GitHubCopilotLSP {
-    func sendRequest<E: GitHubCopilotRequestType>(_ endpoint: E) async throws -> E.Response {
-        try await sendRequest(endpoint.request)
+    func sendRequest<E: GitHubCopilotRequestType>(
+        _ endpoint: E,
+        timeout: TimeInterval? = nil
+    ) async throws -> E.Response {
+        if let timeout {
+            return try await withCheckedThrowingContinuation { continuation in
+                self.sendRequest(endpoint.request, timeout: timeout) { result in
+                    continuation.resume(with: result)
+                }
+            }
+        } else {
+            return try await sendRequest(endpoint.request)
+        }
     }
 }
 

Tool/Sources/GitHubCopilotService/Services/GitHubCopilotChatService.swift

@@ -45,11 +45,13 @@ public final class GitHubCopilotChatService: BuiltinExtensionChatServiceType {
             workspaceFolder: workspace.projectURL.path
         ))
 
-        var cont: AsyncThrowingStream<String, Error>.Continuation? = nil
         let stream = AsyncThrowingStream<String, Error> { continuation in
-            cont = continuation
             let startTimestamp = Date()
 
+            continuation.onTermination = { _ in
+                Task { service.unregisterNotificationHandler(id: id) }
+            }
+
             service.registerNotificationHandler(id: id) { notification, data in
                 // just incase the conversation is stuck, we will cancel it after timeout
                 if Date().timeIntervalSince(startTimestamp) > 60 * 30 {
@@ -71,11 +73,9 @@ public final class GitHubCopilotChatService: BuiltinExtensionChatServiceType {
                             if let error = progress.value.error,
                                progress.value.cancellationReason == nil
                             {
-                                continuation.finish(throwing: ServerError.serverError(
-                                    code: 0,
-                                    message: error,
-                                    data: nil
-                                ))
+                                continuation.finish(
+                                    throwing: GitHubCopilotError.chatEndsWithError(error)
+                                )
                             } else {
                                 continuation.finish()
                             }
@@ -99,22 +99,22 @@ public final class GitHubCopilotChatService: BuiltinExtensionChatServiceType {
                     return false
                 }
             }
-        }
 
-        do {
-            let createResponse = try await service.server.sendRequest(request)
-            cont?.onTermination = { _ in
-                Task {
-                    service.unregisterNotificationHandler(id: id)
+            Task {
+                do {
+                    // this will return when the response is generated.
+                    let createResponse = try await service.server.sendRequest(request, timeout: 120)
                     _ = try await service.server.sendRequest(
                         GitHubCopilotRequest.ConversationDestroy(requestBody: .init(
                             conversationId: createResponse.conversationId
                         ))
                     )
+                } catch let error as ServerError {
+                    continuation.finish(throwing: GitHubCopilotError.languageServerError(error))
+                } catch {
+                    continuation.finish(throwing: error)
                 }
             }
-        } catch {
-            cont?.finish(throwing: error)
         }
 
         return stream
@@ -168,7 +168,7 @@ extension GitHubCopilotChatService {
     func createNewMessage(references: [RetrievedContent], message: String) -> String {
         return message
     }
-    
+
     struct JSONRPC<Params: Decodable>: Decodable {
         var jsonrpc: String
         var method: String