feat: honor $/cancelRequest by cancelling the in-flight task

pyflowlauncher/launcher.py

@@ -4,7 +4,7 @@
 import os
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Awaitable, Callable, Optional
+from typing import Any, Awaitable, Callable, Dict, Optional
 
 from .api import Api
 from .base import pyFlowLauncherObject
@@ -95,6 +95,7 @@ async def _fuzzy_search(self, query: str, text: str) -> MatchData:
 
     async def run(self, dispatch: Callable[[str, list], Awaitable[Any]]) -> None:
         tasks: set = set()
+        in_flight: Dict[Any, asyncio.Task] = {}
         async for request in self._client.messages():
             request_id = request.get('id')
             method = request.get('method', '')
@@ -110,6 +111,17 @@ async def run(self, dispatch: Callable[[str, list], Awaitable[Any]]) -> None:
                 self._respond(request_id, {})
                 continue
 
+            if method == '$/cancelRequest':
+                # StreamJsonRpc cancels a superseded request (e.g. the user kept
+                # typing) with params {'id': <request id>}. Cancel the in-flight
+                # task so slow queries stop doing work nobody will see.
+                cancel_params = request.get('params')
+                cancel_id = cancel_params.get('id') if isinstance(cancel_params, dict) else None
+                cancelled = in_flight.get(cancel_id)
+                if cancelled is not None:
+                    cancelled.cancel()
+                continue
+
             if method.startswith('$/'):
                 continue
 
@@ -132,7 +144,14 @@ async def run(self, dispatch: Callable[[str, list], Awaitable[Any]]) -> None:
                 self._handle_request(request_id, method, params, dispatch)
             )
             tasks.add(task)
-            task.add_done_callback(tasks.discard)
+            if request_id is not None:
+                in_flight[request_id] = task
+
+            def _cleanup(done: asyncio.Task, rid: Any = request_id) -> None:
+                tasks.discard(done)
+                if in_flight.get(rid) is done:
+                    del in_flight[rid]
+            task.add_done_callback(_cleanup)
 
         if tasks:
             await asyncio.gather(*tasks, return_exceptions=True)
@@ -146,6 +165,13 @@ async def _handle_request(
     ) -> None:
         try:
             result = await dispatch(method, params)
+        except asyncio.CancelledError:
+            # StreamJsonRpc's RequestCanceled error; the host discards it, but
+            # answering every request keeps the envelope contract uniform.
+            self._client.send({'id': request_id, 'result': None, 'error': {
+                'code': -32800, 'message': 'Request cancelled',
+            }})
+            raise
         except Exception:
             self.logger.exception("Unhandled error dispatching %r", method)
             self._respond(request_id, {

tests/integration/test_v2_protocol.py

@@ -203,6 +203,50 @@ async def dispatch(method: str, params: list) -> Any:
         assert 'query' not in dispatched
 
 
+class TestV2Cancellation:
+
+    def test_cancel_request_cancels_in_flight_query(self):
+        """$/cancelRequest must stop a slow query instead of letting it finish."""
+        plugin = Plugin(launcher=FlowLauncherV2())
+        completed = []
+
+        @plugin.on_method
+        async def query(q: str):
+            await asyncio.sleep(5)
+            completed.append(q)
+            return Result(title="too late")
+
+        responses = run(plugin, [
+            {'id': 1, 'method': 'query', 'params': [{'search': 'slow'}, {}]},
+            {'method': '$/cancelRequest', 'params': {'id': 1}},
+            {'id': 2, 'method': 'close', 'params': []},
+        ])
+        assert completed == []
+        resp = query_response(responses, 1)
+        assert resp['result'] is None
+        assert resp['error']['code'] == -32800
+
+    def test_cancel_unknown_id_is_ignored(self):
+        plugin = make_plugin()
+        responses = run(plugin, [
+            {'method': '$/cancelRequest', 'params': {'id': 99}},
+            {'id': 1, 'method': 'query', 'params': [{'search': 'World'}, {}]},
+            {'id': 2, 'method': 'close', 'params': []},
+        ])
+        assert query_response(responses, 1)['result']['result'][0]['Title'] == 'Hello, World!'
+        assert all(r.get('id') != 99 for r in responses)
+
+    def test_cancel_with_malformed_params_is_ignored(self):
+        plugin = make_plugin()
+        responses = run(plugin, [
+            {'method': '$/cancelRequest'},
+            {'method': '$/cancelRequest', 'params': [1]},
+            {'id': 1, 'method': 'query', 'params': [{'search': 'World'}, {}]},
+            {'id': 2, 'method': 'close', 'params': []},
+        ])
+        assert query_response(responses, 1)['result']['result'][0]['Title'] == 'Hello, World!'
+
+
 class TestV2Notifications:
 
     def test_cancel_request_silently_ignored(self):