supports `$/cancelRequest`

4 files changed, 14 insertions, 4 deletions

diff --git a/changelog.md b/changelog.md
index 486b0654..df0a8350 100644
--- a/changelog.md
+++ b/changelog.md
@@ -9,6 +9,7 @@
 * `NEW` diagnostics:
   + `different-requires`
 * `NEW` supports `---@CustomClass<string, number>`
+* `NEW` supports `$/cancelRequest`
 * `CHG` hover: improve showing multi defines
 * `CHG` hover: improve showing multi comments at enums
 * `CHG` hint: `Lua.hint.paramName` now supports `Disable`, `Literal` and `All`
diff --git a/script/await.lua b/script/await.lua
index 2d27ae2c..495f381c 100644
--- a/script/await.lua
+++ b/script/await.lua
@@ -1,11 +1,12 @@
 local timer = require 'timer'
-local util  = require 'utility'
+
+local wkmt = { __mode = 'k' }
 
 ---@class await
 local m = {}
 m.type = 'await'
 
-m.coMap = setmetatable({}, { __mode = 'k' })
+m.coMap = setmetatable({}, wkmt)
 m.idMap = {}
 m.delayQueue = {}
 m.delayQueueIndex = 1
@@ -74,7 +75,7 @@ function m.setID(id, co)
         return
     end
     if not m.idMap[id] then
-        m.idMap[id] = setmetatable({}, { __mode = 'k' })
+        m.idMap[id] = setmetatable({}, wkmt)
     end
     m.idMap[id][co] = true
 end
@@ -85,6 +86,7 @@ function m.close(id)
     if not map then
         return
     end
+    m.idMap[id] = nil
     for co in pairs(map) do
         if coroutine.status(co) == 'suspended' then
             map[co] = nil
diff --git a/script/proto/proto.lua b/script/proto/proto.lua
index a0880011..61306b9f 100644
--- a/script/proto/proto.lua
+++ b/script/proto/proto.lua
@@ -4,7 +4,6 @@ local await      = require 'await'
 local pub        = require 'pub'
 local jsonrpc    = require 'jsonrpc'
 local define     = require 'proto.define'
-local timer      = require 'timer'
 local json       = require 'json'
 
 local reqCounter = util.counter()
@@ -121,6 +120,9 @@ function m.doMethod(proto)
     end
     await.call(function ()
         --log.debug('Start method:', method)
+        if proto.id then
+            await.setID('proto:' .. proto.id)
+        end
         local clock = os.clock()
         local ok = true
         local res
@@ -134,6 +136,7 @@ function m.doMethod(proto)
             if not proto.id then
                 return
             end
+            await.close('proto:' .. proto.id)
             if ok then
                 m.response(proto.id, res)
             else
diff --git a/script/provider/provider.lua b/script/provider/provider.lua
index 2b6a1f7c..070652ef 100644
--- a/script/provider/provider.lua
+++ b/script/provider/provider.lua
@@ -805,6 +805,10 @@ proto.on('textDocument/onTypeFormatting', function (params)
     return results
 end)
 
+proto.on('$/cancelRequest', function (params)
+    await.close('proto:' .. params.id)
+end)
+
 proto.on('$/requestHint', function (params)
     local core = require 'core.hint'
     if not config.get 'Lua.hint.enable' then