import asyncio import importlib import os import sys import threading from pathlib import Path from fastapi.testclient import TestClient os.environ["JOB_BACKEND"] = "memory" BACKEND_DIR = Path(__file__).resolve().parents[1] if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) import job_handlers # type: ignore import job_system # type: ignore import risk_control # type: ignore import session_store # type: ignore import audit_store # type: ignore main = importlib.import_module("main") API_KEY_HEADERS = {"X-API-Key": "your-secret-key-here"} def setup_function(): job_system.reset_job_manager() risk_control.reset_risk_controller() session_store.reset_session_store() audit_store.reset_audit_store() main._handlers_registered = False def _completion_payload(): return { "prefix": "hello", "suffix": "", "languageId": "markdown", "model_thinking": "low", "privacy_mode": True, } def test_cancel_endpoint_cancels_running_task(monkeypatch): started = threading.Event() cancelled = threading.Event() async def fake_call_ollama(*args, **kwargs): started.set() try: while True: await asyncio.sleep(0.05) except asyncio.CancelledError: cancelled.set() raise monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama) request_id = "req-cancel-1" with TestClient(main.app) as client: response_box = {} def send_completion(): with client.stream( "POST", "/v1/completions", headers={**API_KEY_HEADERS, "X-Request-Id": request_id}, json=_completion_payload(), ) as response: response_box["status_code"] = response.status_code response_box["body"] = "".join(response.iter_text()) completion_thread = threading.Thread(target=send_completion, daemon=True) completion_thread.start() assert started.wait(timeout=2.0) cancel_response = client.post( "/v1/completions/cancel", headers=API_KEY_HEADERS, json={"request_id": request_id, "reason": "superseded"}, ) assert cancel_response.status_code == 200 assert cancel_response.json() == {"cancelled": True, "status": "ok"} completion_thread.join(timeout=5.0) assert not completion_thread.is_alive() assert cancelled.wait(timeout=2.0) assert "event: cancelled" in response_box["body"] def test_cancel_not_found(): with TestClient(main.app) as client: response = client.post( "/v1/completions/cancel", headers=API_KEY_HEADERS, json={"request_id": "missing", "reason": "abort"}, ) assert response.status_code == 200 assert response.json() == {"cancelled": False, "status": "not_found"}