llm-in-text/backend/tests/test_main_cancel.py

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


class FakeRedis:
    def __init__(self):
        self.acks = []

    async def xack(self, *args):
        self.acks.append(args)

    async def hincrby(self, key, field, amount):
        return 0


class FakeManager:
    def __init__(self):
        self.redis = FakeRedis()
        self.statuses = {}

    async def get_status(self, job_id):
        return self.statuses.get(job_id)

    async def _set_state(self, job_id, state):
        self.statuses[job_id] = state

    async def _metrics(self, job_type):
        return {"queued_count": 0, "running_count": 0}

    async def _emit_event(self, job_id, event, data):
        self.statuses[job_id]["event"] = event

    def _metrics_key(self, job_type):
        return f"metrics:{job_type}"

    def _state_key(self, job_id):
        return f"state:{job_id}"


async def _run_cancelled_after_handler(manager, job_type):
    worker = job_system.RedisWorker(manager)
    await worker._run_message(
        job_type,
        "queue",
        "group",
        "msg-1",
        {"job_id": "job-1"},
        asyncio.Semaphore(1),
    )


def test_redis_worker_acks_when_handler_returns_cancelled_state():
    async def handler(payload, emit, is_cancelled):
        return {"ok": True}

    async def coro():
        manager = FakeManager()
        manager.handlers = {"completion": handler}
        manager.statuses["job-1"] = {
            "request_id": "req-1",
            "type": "completion",
            "status": "running",
            "created_at": 1,
        }
        await _run_cancelled_after_handler(manager, "completion")
        assert manager.redis.acks == [("queue", "group", "msg-1")]

    asyncio.run(coro())


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