diff --git a/pyflowlauncher/launcher.py b/pyflowlauncher/launcher.py index c7b4254..dbf78f8 100644 --- a/pyflowlauncher/launcher.py +++ b/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': }. 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, { diff --git a/tests/integration/test_v2_protocol.py b/tests/integration/test_v2_protocol.py index 5a4369e..1bfe211 100644 --- a/tests/integration/test_v2_protocol.py +++ b/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):