From f0d4ccbeb338f4726fa26abbb4a517454005bad4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?= <96022931+mustafalpyilmaz@users.noreply.github.com> Date: Tue, 3 Mar 2026 13:38:52 +0300 Subject: [PATCH] feat: add POST /v1/cancel/{command_id} endpoint (#1579) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Adds explicit `POST /v1/cancel/{command_id}` REST endpoint to cancel in-flight text and image generation commands by ID - Previously, cancellation only worked via HTTP disconnect (client closes SSE connection, triggering `anyio.get_cancelled_exc_class()`). Non-streaming clients and external consumers had no way to cleanly cancel active generation - The endpoint looks up the command in active generation queues, sends `TaskCancelled` to notify workers, and closes the sender stream. Returns 404 (OpenAI error format) if the command is not found or already completed ## Changes | File | Change | |------|--------| | `src/exo/shared/types/api.py` | Add `CancelCommandResponse` model | | `src/exo/master/api.py` | Import, route registration, `cancel_command` handler | | `src/exo/master/tests/test_cancel_command.py` | 3 test cases (404, text cancel, image cancel) | | `docs/api.md` | Document new endpoint + update summary table | ## Design Decisions - Uses `self._send()` (not raw `command_sender.send()`) — respects API pause state during elections - Uses `raise HTTPException` — feeds into exo's centralized OpenAI-style error handler - Returns typed `CancelCommandResponse` — consistent with `CreateInstanceResponse` / `DeleteInstanceResponse` patterns - `sender.close()` is idempotent so concurrent cancel requests for the same command are harmless ## Test Plan - [x] `test_cancel_nonexistent_command_returns_404` — verifies 404 with OpenAI error format - [x] `test_cancel_active_text_generation` — verifies 200, `sender.close()` called, `TaskCancelled` sent with correct `cancelled_command_id` - [x] `test_cancel_active_image_generation` — same verification for image queue - [x] `basedpyright` — 0 new errors - [x] `ruff check` — all checks passed - [x] `pytest` — 3/3 new tests pass, no regressions --------- Co-authored-by: Evan --- docs/api.md | 22 ++++++ src/exo/master/api.py | 21 ++++++ src/exo/master/tests/test_cancel_command.py | 77 +++++++++++++++++++++ src/exo/shared/types/api.py | 5 ++ 4 files changed, 125 insertions(+) create mode 100644 src/exo/master/tests/test_cancel_command.py diff --git a/docs/api.md b/docs/api.md index 9c4ec096..c25e3bc4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -183,6 +183,27 @@ Same schema as `/v1/chat/completions`. **Response:** Chat completion plus benchmarking metrics. +### Cancel Command + +**POST** `/v1/cancel/{command_id}` + +Cancels an active generation command (text or image). Notifies workers and closes the stream. + +**Path parameters:** + +* `command_id`: string, ID of the command to cancel + +**Response (example):** + +```json +{ + "message": "Command cancelled.", + "command_id": "cmd-abc-123" +} +``` + +Returns 404 if the command is not found or already completed. + ## 5. Image Generation & Editing ### Image Generation @@ -266,6 +287,7 @@ GET /v1/models POST /v1/chat/completions POST /bench/chat/completions +POST /v1/cancel/{command_id} POST /v1/images/generations POST /bench/images/generations diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 60a7311f..5da82662 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -73,6 +73,7 @@ from exo.shared.types.api import ( BenchChatCompletionResponse, BenchImageGenerationResponse, BenchImageGenerationTaskParams, + CancelCommandResponse, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest, @@ -322,6 +323,7 @@ class API: self.app.get("/images/{image_id}")(self.get_image) self.app.post("/v1/messages", response_model=None)(self.claude_messages) self.app.post("/v1/responses", response_model=None)(self.openai_responses) + self.app.post("/v1/cancel/{command_id}")(self.cancel_command) # Ollama API self.app.head("/ollama/")(self.ollama_version) @@ -567,6 +569,25 @@ class API: instance_id=instance_id, ) + async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse: + """Cancel an active command by closing its stream and notifying workers.""" + sender = self._text_generation_queues.get( + command_id + ) or self._image_generation_queues.get(command_id) + if sender is None: + raise HTTPException( + status_code=404, + detail="Command not found or already completed", + ) + + await self._send(TaskCancelled(cancelled_command_id=command_id)) + sender.close() + + return CancelCommandResponse( + message="Command cancelled.", + command_id=command_id, + ) + async def _token_chunk_stream( self, command_id: CommandId ) -> AsyncGenerator[ diff --git a/src/exo/master/tests/test_cancel_command.py b/src/exo/master/tests/test_cancel_command.py new file mode 100644 index 00000000..b64fc0cc --- /dev/null +++ b/src/exo/master/tests/test_cancel_command.py @@ -0,0 +1,77 @@ +# pyright: reportUnusedFunction=false, reportAny=false +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from exo.shared.types.common import CommandId + + +def _make_api() -> Any: + """Create a minimal API instance with cancel route and error handler.""" + from exo.master.api import API + + app = FastAPI() + api = object.__new__(API) + api.app = app + api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage] + api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage] + api._send = AsyncMock() # pyright: ignore[reportPrivateUsage] + api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage] + app.post("/v1/cancel/{command_id}")(api.cancel_command) + return api + + +def test_cancel_nonexistent_command_returns_404() -> None: + """Cancel for an unknown command_id returns 404 in OpenAI error format.""" + api = _make_api() + client = TestClient(api.app) + + response = client.post("/v1/cancel/nonexistent-id") + assert response.status_code == 404 + data: dict[str, Any] = response.json() + assert "error" in data + assert data["error"]["message"] == "Command not found or already completed" + assert data["error"]["type"] == "Not Found" + assert data["error"]["code"] == 404 + + +def test_cancel_active_text_generation() -> None: + """Cancel an active text generation command: returns 200, sender.close() called.""" + api = _make_api() + client = TestClient(api.app) + + cid = CommandId("text-cmd-123") + sender = MagicMock() + api._text_generation_queues[cid] = sender + + response = client.post(f"/v1/cancel/{cid}") + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["message"] == "Command cancelled." + assert data["command_id"] == str(cid) + sender.close.assert_called_once() + api._send.assert_called_once() + task_cancelled = api._send.call_args[0][0] + assert task_cancelled.cancelled_command_id == cid + + +def test_cancel_active_image_generation() -> None: + """Cancel an active image generation command: returns 200, sender.close() called.""" + api = _make_api() + client = TestClient(api.app) + + cid = CommandId("img-cmd-456") + sender = MagicMock() + api._image_generation_queues[cid] = sender + + response = client.post(f"/v1/cancel/{cid}") + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["message"] == "Command cancelled." + assert data["command_id"] == str(cid) + sender.close.assert_called_once() + api._send.assert_called_once() + task_cancelled = api._send.call_args[0][0] + assert task_cancelled.cancelled_command_id == cid diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 034414a5..511bc754 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -262,6 +262,11 @@ class DeleteInstanceResponse(BaseModel): instance_id: InstanceId +class CancelCommandResponse(BaseModel): + message: str + command_id: CommandId + + ImageSize = Literal[ "auto", "512x512",