Add abort request processing

ncoghlan · ncoghlan · commit e538f2ae495b · 2025-07-25T04:19:26.000+10:00
diff --git a/examples/plugins/prompt-prefix/src/plugin.py b/examples/plugins/prompt-prefix/src/plugin.py
@@ -39,12 +39,12 @@ async def preprocess_prompt(
     message: UserMessage,
 ) -> UserMessageDict | None:
     """Naming the function 'preprocess_prompt' implicitly registers it."""
-    print(f"Running prompt preprocessor hook from {__file__} with {ctl.plugin_config}")
     if ctl.global_config.enable_inplace_status_demo:
         # Run an in-place status prompt update demonstration
         status_block = await ctl.notify_start("Starting task (shows a static icon).")
         status_updates = (
             (status_block.notify_working, "Task in progress (shows a dynamic icon)."),
+            (status_block.notify_waiting, "Task is blocked (shows a static icon)."),
             (status_block.notify_error, "Reporting an error status."),
             (status_block.notify_canceled, "Reporting cancellation."),
             (
@@ -55,9 +55,11 @@ async def preprocess_prompt(
         status_duration = ctl.global_config.inplace_status_duration / len(
             status_updates
         )
-        for notification, status_text in status_updates:
-            await asyncio.sleep(status_duration)
-            await notification(status_text)
+        async with status_block.notify_aborted("Task genuinely cancelled."):
+            for notification, status_text in status_updates:
+                await asyncio.sleep(status_duration)
+                await notification(status_text)
+
     modified_message = message.to_dict()
     # Add a prefix to all user messages
     prefix_text = ctl.plugin_config.prefix
diff --git a/sdk-schema/sync-sdk-schema.py b/sdk-schema/sync-sdk-schema.py
@@ -363,9 +363,11 @@ def _infer_schema_unions() -> None:
     "LlmChannelPredictCreationParameterDict": "PredictionChannelRequestDict",
     "RepositoryChannelDownloadModelCreationParameter": "DownloadModelChannelRequest",
     "RepositoryChannelDownloadModelCreationParameterDict": "DownloadModelChannelRequestDict",
-    # Prettier plugin channel type names
+    # Prettier plugin channel message names
     "PluginsChannelSetPromptPreprocessorToClientPacketPreprocess": "PromptPreprocessingRequest",
     "PluginsChannelSetPromptPreprocessorToClientPacketPreprocessDict": "PromptPreprocessingRequestDict",
+    "PluginsChannelSetPromptPreprocessorToServerPacketAborted": "PromptPreprocessingAborted",
+    "PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict": "PromptPreprocessingAbortedDict",
     "PluginsChannelSetPromptPreprocessorToServerPacketComplete": "PromptPreprocessingComplete",
     "PluginsChannelSetPromptPreprocessorToServerPacketCompleteDict": "PromptPreprocessingCompleteDict",
     "PluginsChannelSetPromptPreprocessorToServerPacketError": "PromptPreprocessingError",
diff --git a/src/lmstudio/_sdk_models/__init__.py b/src/lmstudio/_sdk_models/__init__.py
@@ -420,8 +420,6 @@
     "PluginsChannelSetPredictionLoopHandlerToServerPacketErrorDict",
     "PluginsChannelSetPromptPreprocessorToClientPacketAbort",
     "PluginsChannelSetPromptPreprocessorToClientPacketAbortDict",
-    "PluginsChannelSetPromptPreprocessorToServerPacketAborted",
-    "PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict",
     "PluginsChannelSetToolsProviderToClientPacketAbortToolCall",
     "PluginsChannelSetToolsProviderToClientPacketAbortToolCallDict",
     "PluginsChannelSetToolsProviderToClientPacketCallTool",
@@ -522,6 +520,8 @@
     "ProcessingUpdateToolStatusCreateDict",
     "ProcessingUpdateToolStatusUpdate",
     "ProcessingUpdateToolStatusUpdateDict",
+    "PromptPreprocessingAborted",
+    "PromptPreprocessingAbortedDict",
     "PromptPreprocessingComplete",
     "PromptPreprocessingCompleteDict",
     "PromptPreprocessingError",
@@ -5061,8 +5061,8 @@ class PluginsChannelSetPromptPreprocessorToClientPacketAbortDict(TypedDict):
     taskId: str
 
 
-class PluginsChannelSetPromptPreprocessorToServerPacketAborted(
-    LMStudioStruct["PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict"],
+class PromptPreprocessingAborted(
+    LMStudioStruct["PromptPreprocessingAbortedDict"],
     kw_only=True,
     tag_field="type",
     tag="aborted",
@@ -5071,7 +5071,7 @@ class PluginsChannelSetPromptPreprocessorToServerPacketAborted(
     task_id: str = field(name="taskId")
 
 
-class PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict(TypedDict):
+class PromptPreprocessingAbortedDict(TypedDict):
     """Corresponding typed dictionary definition for PluginsChannelSetPromptPreprocessorToServerPacketAborted.
 
     NOTE: Multi-word keys are defined using their camelCase form,
@@ -10199,14 +10199,12 @@ class PseudoPluginsChannelRegisterDevelopmentPluginDict(TypedDict):
     | PluginsChannelSetPromptPreprocessorToClientPacketAbortDict
 )
 PluginsChannelSetPromptPreprocessorToServerPacket = (
-    PromptPreprocessingComplete
-    | PluginsChannelSetPromptPreprocessorToServerPacketAborted
-    | PromptPreprocessingError
+    PromptPreprocessingComplete | PromptPreprocessingAborted | PromptPreprocessingError
 )
 PluginsChannelSetPromptPreprocessorToServerPacketDict = (
     PromptPreprocessingErrorDict
     | PromptPreprocessingCompleteDict
-    | PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict
+    | PromptPreprocessingAbortedDict
 )
 
 
diff --git a/src/lmstudio/_ws_impl.py b/src/lmstudio/_ws_impl.py
@@ -290,7 +290,7 @@ async def _log_thread_execution(self) -> None:
         try:
             # Run the event loop until termination is requested
             await never_set.wait()
-        except asyncio.CancelledError:
+        except (asyncio.CancelledError, GeneratorExit):
             raise
         except BaseException:
             err_msg = "Terminating websocket thread due to exception"
@@ -359,7 +359,7 @@ async def _logged_ws_handler(self) -> None:
         self._logger.info("Websocket handling task started")
         try:
             await self._handle_ws()
-        except asyncio.CancelledError:
+        except (asyncio.CancelledError, GeneratorExit):
             raise
         except BaseException:
             err_msg = "Terminating websocket task due to exception"
diff --git a/src/lmstudio/plugin/__init__.py b/src/lmstudio/plugin/__init__.py
@@ -33,7 +33,7 @@
 # Prompt preprocessing hook
 # * [DONE] emit a status notification block when the demo plugin fires
 # * [DONE] add a global plugin config to control the in-place status update demo
-# * handle "Abort" requests from server (including sending "Aborted" responses)
+# * [DONE] handle "Abort" requests from server (including sending "Aborted" responses)
 # * [DONE] catch hook invocation failures and send "Error" responses
 # * [DONE] this includes adding runtime checks for the hook returning the wrong type
 #
diff --git a/src/lmstudio/plugin/hooks/common.py b/src/lmstudio/plugin/hooks/common.py
@@ -1,17 +1,23 @@
 """Common utilities to invoke and support plugin hook implementations."""
 
+import asyncio
+
+from contextlib import asynccontextmanager
 from datetime import datetime, UTC
 from pathlib import Path
 from random import randrange
 from typing import (
     Any,
+    AsyncIterator,
     Awaitable,
     Callable,
     Generic,
     TypeAlias,
     TypeVar,
 )
 
+from anyio import move_on_after
+
 from ...async_api import AsyncSession
 from ...schemas import DictObject
 from ..._sdk_models import (
@@ -48,6 +54,10 @@ class AsyncSessionPlugins(AsyncSession):
 SendMessageCallback: TypeAlias = Callable[[DictObject], Awaitable[Any]]
 
 
+class ServerRequestError(RuntimeError):
+    """Plugin received an invalid request from the API server."""
+
+
 class HookController(Generic[TPluginRequest, TPluginConfigSchema, TGlobalConfigSchema]):
     """Common base class for plugin hook API access controllers."""
 
@@ -117,3 +127,14 @@ async def notify_canceled(self, message: str) -> None:
     async def notify_done(self, message: str) -> None:
         """Report task completion in the status block."""
         await self._update_ui(self._id, "done", message)
+
+    @asynccontextmanager
+    async def notify_aborted(self, message: str) -> AsyncIterator[None]:
+        """Report asyncio.CancelledError as cancellation in the status block."""
+        try:
+            yield
+        except asyncio.CancelledError:
+            # Allow the notification to be sent, but don't necessarily wait for the reply
+            with move_on_after(0.2, shield=True):
+                await self.notify_canceled(message)
+            raise
diff --git a/src/lmstudio/plugin/hooks/prompt_preprocessor.py b/src/lmstudio/plugin/hooks/prompt_preprocessor.py
@@ -1,9 +1,13 @@
 """Invoking and supporting prompt preprocessor hook implementations."""
 
+import asyncio
+
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from traceback import format_tb
 from typing import (
     Any,
+    AsyncIterator,
     Awaitable,
     Callable,
     Generic,
@@ -13,6 +17,7 @@
 )
 
 from anyio import create_task_group
+from anyio.abc import TaskGroup
 
 from ..._logging import new_logger
 from ...schemas import DictObject, EmptyDict, ValidationError
@@ -29,6 +34,7 @@
     ProcessingUpdate,
     ProcessingUpdateStatusCreate,
     ProcessingUpdateStatusUpdate,
+    PromptPreprocessingAbortedDict,
     PromptPreprocessingCompleteDict,
     PromptPreprocessingErrorDict,
     PromptPreprocessingRequest,
@@ -41,12 +47,12 @@
     AsyncSessionPlugins,
     HookController,
     SendMessageCallback,
+    ServerRequestError,
     StatusBlockController,
     TPluginConfigSchema,
     TGlobalConfigSchema,
 )
 
-
 # Available as lmstudio.plugin.hooks.*
 __all__ = [
     "PromptPreprocessorController",
@@ -88,7 +94,7 @@ def iter_message_events(
             case None:
                 # Server can only terminate the link by closing the websocket
                 pass
-            case {"type": "abort", "task_id": str(task_id)}:
+            case {"type": "abort", "taskId": str(task_id)}:
                 yield PromptPreprocessingAbortEvent(task_id)
             case {"type": "preprocess"} as request_dict:
                 parsed_request = PromptPreprocessingRequest._from_any_api_dict(
@@ -101,10 +107,10 @@ def iter_message_events(
     def handle_rx_event(self, event: PromptPreprocessingRxEvent) -> None:
         match event:
             case PromptPreprocessingAbortEvent(task_id):
-                self._logger.info(f"Aborting {task_id}", task_id=task_id)
+                self._logger.debug(f"Aborting {task_id}", task_id=task_id)
             case PromptPreprocessingRequestEvent(request):
                 task_id = request.task_id
-                self._logger.info(
+                self._logger.debug(
                     "Received prompt preprocessing request", task_id=task_id
                 )
             case ChannelFinishedEvent(_):
@@ -198,16 +204,18 @@ class PromptPreprocessor(Generic[TPluginConfigSchema, TGlobalConfigSchema]):
     def __post_init__(self) -> None:
         self._logger = logger = new_logger(__name__)
         logger.update_context(plugin_name=self.plugin_name)
+        self._abort_events: dict[str, asyncio.Event] = {}
 
     async def process_requests(
         self, session: AsyncSessionPlugins, notify_ready: Callable[[], Any]
     ) -> None:
+        """Create plugin channel and wait for server requests."""
         logger = self._logger
         endpoint = PromptPreprocessingEndpoint()
         async with session._create_channel(endpoint) as channel:
             notify_ready()
             logger.info("Opened channel to receive prompt preprocessing requests...")
-            send_cb = channel.send_message
+            send_message = channel.send_message
             async with create_task_group() as tg:
                 logger.debug("Waiting for prompt preprocessing requests...")
                 async for contents in channel.rx_stream():
@@ -218,6 +226,10 @@ async def process_requests(
                         logger.debug("Handling prompt preprocessing channel event")
                         endpoint.handle_rx_event(event)
                         match event:
+                            case PromptPreprocessingAbortEvent():
+                                await self._abort_hook_invocation(
+                                    event.arg, send_message
+                                )
                             case PromptPreprocessingRequestEvent():
                                 logger.debug(
                                     "Running prompt preprocessing request hook"
@@ -228,28 +240,76 @@ async def process_requests(
                                     self.plugin_config_schema,
                                     self.global_config_schema,
                                 )
-                                tg.start_soon(self._invoke_hook, ctl, send_cb)
+                                tg.start_soon(self._invoke_hook, ctl, send_message)
                     if endpoint.is_finished:
                         break
 
+    async def _abort_hook_invocation(
+        self, task_id: str, send_response: SendMessageCallback
+    ) -> None:
+        """Abort the specified hook invocation (if it is still running)."""
+        abort_event = self._abort_events.get(task_id, None)
+        if abort_event is not None:
+            abort_event.set()
+        response = PromptPreprocessingAbortedDict(
+            type="aborted",
+            taskId=task_id,
+        )
+        await send_response(response)
+
+    async def _cancel_on_event(
+        self, tg: TaskGroup, event: asyncio.Event, message: str
+    ) -> None:
+        await event.wait()
+        self._logger.info(message)
+        tg.cancel_scope.cancel()
+
+    @asynccontextmanager
+    async def _registered_hook_invocation(
+        self, task_id: str
+    ) -> AsyncIterator[asyncio.Event]:
+        logger = self._logger
+        abort_events = self._abort_events
+        if task_id in abort_events:
+            err_msg = f"Hook invocation already in progress for {task_id}"
+            raise ServerRequestError(err_msg)
+        abort_events[task_id] = abort_event = asyncio.Event()
+        try:
+            async with create_task_group() as tg:
+                tg.start_soon(
+                    self._cancel_on_event,
+                    tg,
+                    abort_event,
+                    f"Aborting request {task_id}",
+                )
+                logger.info(f"Processing request {task_id}")
+                yield abort_event
+                tg.cancel_scope.cancel()
+        finally:
+            abort_events.pop(task_id, None)
+        if abort_event.is_set():
+            completion_message = f"Aborted request {task_id}"
+        else:
+            completion_message = f"Processed request {task_id}"
+        logger.info(completion_message)
+
     async def _invoke_hook(
         self,
         ctl: PromptPreprocessorController[TPluginConfigSchema, TGlobalConfigSchema],
         send_response: SendMessageCallback,
     ) -> None:
         logger = self._logger
-        request = ctl.request
-        message = request.input
-        expected_cls = UserMessage
-        if not isinstance(message, expected_cls):
-            logger.error(
-                f"Received {type(message).__name__!r} ({expected_cls.__name__!r} expected)"
-            )
-            return
+        task_id = ctl.task_id
+        message = ctl.request.input
         error_details: SerializedLMSExtendedErrorDict | None = None
         response_dict: UserMessageDict
+        expected_cls = UserMessage
         try:
-            response = await self.hook_impl(ctl, message)
+            if not isinstance(message, expected_cls):
+                err_msg = f"Received {type(message).__name__!r} ({expected_cls.__name__!r} expected)"
+                raise ServerRequestError(err_msg)
+            async with self._registered_hook_invocation(task_id) as abort_event:
+                response = await self.hook_impl(ctl, message)
         except Exception as exc:
             err_msg = "Error calling prompt preprocessing hook"
             logger.error(err_msg, exc_info=True, exc=repr(exc))
@@ -259,8 +319,11 @@ async def _invoke_hook(
                 cause=ui_cause, stack="\n".join(format_tb(exc.__traceback__))
             )
         else:
+            if abort_event.is_set():
+                # Processing was aborted by the server, skip sending a response
+                return
             if response is None:
-                # No change to message
+                logger.debug("No changes made to preprocessed prompt")
                 response_dict = message.to_dict()
             else:
                 logger.debug(
@@ -291,13 +354,13 @@ async def _invoke_hook(
             error_details.update(common_error_args)
             channel_message = PromptPreprocessingErrorDict(
                 type="error",
-                taskId=request.task_id,
+                taskId=task_id,
                 error=error_details,
             )
         else:
             channel_message = PromptPreprocessingCompleteDict(
                 type="complete",
-                taskId=request.task_id,
+                taskId=task_id,
                 processed=response_dict,
             )
         await send_response(channel_message)

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@`
`33`	`33`	`# Prompt preprocessing hook`
`34`	`34`	`# * [DONE] emit a status notification block when the demo plugin fires`
`35`	`35`	`# * [DONE] add a global plugin config to control the in-place status update demo`
`36`		`-# * handle "Abort" requests from server (including sending "Aborted" responses)`
	`36`	`+# * [DONE] handle "Abort" requests from server (including sending "Aborted" responses)`
`37`	`37`	`# * [DONE] catch hook invocation failures and send "Error" responses`
`38`	`38`	`# * [DONE] this includes adding runtime checks for the hook returning the wrong type`
`39`	`39`	`#`