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",