cancellation test

Author: keynmol

modules/app/src/main/scala/LSPCancelRequest.scala

@@ -19,13 +19,10 @@ package langoustine.lsp.app
 import jsonrpclib.CallId
 import jsonrpclib.fs2.CancelTemplate
 
-private case class LSPCancelRequest(id: CallId) derives io.circe.Codec.AsObject
-
-object LSPCancelRequest:
-  // import com.github.plokhotnyuk.jsoniter_scala.macros.*
-  // given JsonValueCodec[LSPCancelRequest] = JsonCodecMaker.make[LSPCancelRequest]
-  // given Codec[LSPCancelRequest] = Codec.fromJsonCodec
+private[langoustine] case class LSPCancelRequest(id: CallId)
+    derives io.circe.Codec.AsObject
 
+private[langoustine] object LSPCancelRequest:
   val cancelTemplate: CancelTemplate = CancelTemplate
     .make[LSPCancelRequest](
       "$/cancelRequest",

modules/tests/src/test/scala/LSPTests.scala

@@ -18,13 +18,16 @@ package tests.core
 
 import scala.concurrent.duration.*
 
+import langoustine.testkit.*
+
 import cats.effect.IO
 import jsonrpclib.fs2.catsMonadic
 import jsonrpclib.{fs2 as _, *}
 import langoustine.lsp.*
 import langoustine.lsp.all.*
 import langoustine.lsp.structures.InitializeParams.ClientInfo
 import langoustine.lsp.structures.InitializeResult.ServerInfo
+import cats.syntax.all.*
 
 object LSPTests extends weaver.SimpleIOSuite:
 
@@ -138,6 +141,56 @@ object LSPTests extends weaver.SimpleIOSuite:
 
   }
 
+  serverTest("cancellation") { log =>
+    import requests.*
+
+    val NumParallelCalls = 5
+
+    IO.ref(Set.empty[CallId.StringId])
+      .parProduct(cats.effect.std.CountDownLatch[IO](NumParallelCalls))
+      .flatMap: (cancelled, latch) =>
+        val server = basicServer.handleRequest(textDocument.documentSymbol) {
+          in =>
+            log.info(s"in handler for ${in.params.textDocument.uri}") *>
+              IO.never
+                .onCancel(
+                  log.info(
+                    s"Cancelled ${in.params.textDocument.uri}"
+                  ) *>
+                    cancelled.update(
+                      _ + CallId.StringId(in.params.textDocument.uri.value)
+                    ) *>
+                    latch.release
+                )
+                .as(None)
+        }
+
+        Probe
+          .create(server, log)
+          .use: probe =>
+            val requests = List.tabulate(NumParallelCalls): i =>
+              val callID = CallId.StringId(s"document $i")
+              probe
+                .requestAndForget(
+                  callID,
+                  textDocument.documentSymbol,
+                  DocumentSymbolParams(
+                    TextDocumentIdentifier(
+                      uri = DocumentUri(callID.string)
+                    )
+                  )
+                )
+                .as(callID)
+
+            for
+              callIds      <- requests.parSequence
+              _            <- callIds.parTraverse(probe.cancel)
+              _            <- latch.await.timeout(5.seconds)
+              allCancelled <- cancelled.get
+            yield expect.same(callIds.toSet, allCancelled)
+            end for
+  }
+
   serverTest("documentSymbol") { log =>
     val symbols: Option[Vector[DocumentSymbol]] =
       Option(
@@ -190,51 +243,3 @@ end LSPTests
 
 def basicServer =
   LSPBuilder.create[IO]
-
-//   test("textDocument/documentSymbol") {
-//     import requests.*
-
-//     val symbols: Option[Vector[DocumentSymbol]] =
-//       Opt(
-//         Vector(
-//           DocumentSymbol(
-//             name = "root",
-//             kind = SymbolKind.Array,
-//             range = Range(Position(0, 0), Position(0, 5)),
-//             selectionRange = Range(Position(0, 0), Position(0, 5)),
-//             children = Opt(
-//               Vector(
-//                 DocumentSymbol(
-//                   name = "child",
-//                   kind = SymbolKind.String,
-//                   range = Range(Position(1, 0), Position(1, 5)),
-//                   selectionRange = Range(Position(1, 0), Position(1, 5))
-//                 )
-//               )
-//             )
-//           )
-//         )
-//       )
-
-//     val server = basicServer[Try].handleRequest(textDocument.documentSymbol) {
-//       in =>
-//         Try {
-//           symbols
-//         }
-//     }
-
-//     val (response, _) = request(
-//       server,
-//       CallId.StringId("resp1"),
-//       textDocument.documentSymbol,
-//       DocumentSymbolParams(
-//         TextDocumentIdentifier(DocumentUri("/home/bla.txt"))
-//       )
-//     ).get
-
-//     expect.same(
-//       response,
-//       symbols
-//     )
-//   }
-// end LSPTests

modules/tests/src/test/scala/Probe.scala

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package tests.core
+package langoustine.testkit
 
 import scala.concurrent.duration.*
 
@@ -31,6 +31,8 @@ import jsonrpclib.fs2.catsMonadic
 import jsonrpclib.{fs2 as _, *}
 import langoustine.lsp.*
 import langoustine.lsp.all.*
+import _root_.io.circe.syntax.*
+import langoustine.lsp.app.LSPCancelRequest
 
 class Probe(
     messagesFromServer: SignallingRef[IO, Vector[Message]],
@@ -40,6 +42,13 @@ class Probe(
     timeouts: Timeouts
 ):
 
+  def cancel(callId: CallId) =
+    val msg = InputMessage.NotificationMessage(
+      "$/cancelRequest",
+      Some(Payload(LSPCancelRequest(callId).asJson))
+    )
+    send(msg)
+
   def notification[T <: LSPNotification](t: T, params: t.In): IO[Unit] =
     val msg = InputMessage.NotificationMessage(
       method = t.notificationMethod,
@@ -88,6 +97,20 @@ class Probe(
       .timeout(timeouts.waitForNotifications)
   end waitForNotifications
 
+  def requestAndForget[T <: LSPRequest](
+      callId: CallId,
+      t: T,
+      params: t.In
+  ): IO[Unit] =
+    val msg = InputMessage.RequestMessage(
+      method = t.requestMethod,
+      callId = callId,
+      params = Some(Payload(t.inputToJson.apply(params)))
+    )
+
+    send(msg)
+  end requestAndForget
+
   def request[T <: LSPRequest](callId: CallId, t: T, params: t.In): IO[t.Out] =
     val msg = InputMessage.RequestMessage(
       method = t.requestMethod,
@@ -136,7 +159,7 @@ object Probe:
   ) =
     for
       ch <- FS2Channel
-        .resource[IO]()
+        .resource[IO](cancelTemplate = Some(LSPCancelRequest.cancelTemplate))
         .flatMap(ch => ch.withEndpoints(builder.build(Communicate.channel(ch))))
       _     <- log.info("wut").toResource
       input <- fs2.concurrent.Channel.synchronous[IO, Message].toResource