# -*- coding: utf-8 -*-
# mcpgateway/services/cancellation_service.py
"""Location: ./mcpgateway/services/cancellation_service.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti
Service for tracking and cancelling active tool runs.
Provides a simple in-memory registry for run metadata and an optional async
cancel callback that can be invoked when a cancellation is requested. This
service is intentionally small and designed to be a single-process helper for
local run lifecycle management; the gateway remains authoritative for
cancellation and also broadcasts a `notifications/cancelled` JSON-RPC
notification to connected sessions.
"""
# Future
from __future__ import annotations
# Standard
import asyncio
import json
import time
from typing import Any, Awaitable, Callable, Dict, Optional
# First-Party
from mcpgateway.services.logging_service import LoggingService
from mcpgateway.utils.redis_client import get_redis_client
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)
CancelCallback = Callable[[Optional[str]], Awaitable[None]] # async callback(reason)
class CancellationService:
"""Track active runs and allow cancellation requests.
Note: This is intentionally lightweight — it does not persist state and is
suitable for gateway-local run tracking. The gateway will also broadcast
a `notifications/cancelled` message to connected sessions to inform remote
peers of the cancellation request.
Multi-worker deployments: When Redis is available, cancellation events are
published to the "cancellation:cancel" channel to propagate across workers.
"""
def __init__(self) -> None:
"""Initialize the cancellation service."""
self._runs: Dict[str, Dict[str, Any]] = {}
self._lock = asyncio.Lock()
self._redis = None
self._pubsub_task: Optional[asyncio.Task] = None
self._initialized = False
async def initialize(self) -> None:
"""Initialize Redis pubsub if available for multi-worker support."""
if self._initialized:
return
self._initialized = True
try:
self._redis = await get_redis_client()
if self._redis:
# Start listening for cancellation events from other workers
self._pubsub_task = asyncio.create_task(self._listen_for_cancellations())
logger.info("CancellationService: Redis pubsub initialized for multi-worker cancellation")
except Exception as e:
logger.warning(f"CancellationService: Could not initialize Redis pubsub: {e}")
async def shutdown(self) -> None:
"""Shutdown Redis pubsub listener."""
if self._pubsub_task and not self._pubsub_task.done():
self._pubsub_task.cancel()
try:
await self._pubsub_task
except asyncio.CancelledError:
pass
logger.info("CancellationService: Shutdown complete")
async def _listen_for_cancellations(self) -> None:
"""Listen for cancellation events from other workers via Redis pubsub.
Uses timeout-based polling instead of blocking listen() to allow proper
cancellation handling. This prevents CPU spin loops when the task is cancelled
but stuck waiting on the blocking async iterator.
Raises:
asyncio.CancelledError: When the listener task is cancelled during shutdown.
"""
if not self._redis:
return
pubsub = None
try:
pubsub = self._redis.pubsub()
await pubsub.subscribe("cancellation:cancel")
logger.info("CancellationService: Subscribed to cancellation:cancel channel")
# Use timeout-based polling instead of blocking listen()
# This allows the task to respond to cancellation properly
poll_timeout = 1.0
while True:
try:
message = await asyncio.wait_for(
pubsub.get_message(ignore_subscribe_messages=True, timeout=poll_timeout),
timeout=poll_timeout + 0.5,
)
except asyncio.TimeoutError:
# No message, continue loop to check for cancellation
continue
if message is None:
# Prevent spin if get_message returns None immediately
await asyncio.sleep(0.1)
continue
if message["type"] != "message":
# Sleep on non-message types to prevent spin
await asyncio.sleep(0.1)
continue
try:
data = json.loads(message["data"])
# Normalize run_id to string (handle id=0 which is valid per JSON-RPC)
raw_run_id = data.get("run_id")
run_id = str(raw_run_id) if raw_run_id is not None else None
reason = data.get("reason")
if run_id is not None:
# Cancel locally if we have this run (don't re-publish)
await self._cancel_run_local(run_id, reason=reason)
except Exception as e:
logger.warning(f"Error processing cancellation message: {e}")
except asyncio.CancelledError:
logger.info("CancellationService: Pubsub listener cancelled")
raise
except Exception as e:
logger.error(f"CancellationService: Pubsub listener error: {e}")
finally:
# Clean up pubsub on any exit
if pubsub is not None:
try:
await pubsub.unsubscribe("cancellation:cancel")
try:
await pubsub.aclose()
except AttributeError:
await pubsub.close()
except Exception as e:
logger.debug(f"Error closing cancellation pubsub: {e}")
async def _cancel_run_local(self, run_id: str, reason: Optional[str] = None) -> bool:
"""Cancel a run locally without publishing to Redis (internal use).
Args:
run_id: Unique identifier for the run to cancel.
reason: Optional textual reason for the cancellation request.
Returns:
bool: True if the run was found and cancelled, False if not found.
"""
async with self._lock:
entry = self._runs.get(run_id)
if not entry:
return False
if entry.get("cancelled"):
return True
entry["cancelled"] = True
entry["cancelled_at"] = time.time()
entry["cancel_reason"] = reason
cancel_cb = entry.get("cancel_callback")
logger.info("Tool execution cancelled (from Redis): run_id=%s, reason=%s, tool=%s", run_id, reason or "not specified", entry.get("name", "unknown"))
if cancel_cb:
try:
await cancel_cb(reason)
logger.info("Cancel callback executed for %s", run_id)
except Exception as e:
logger.exception("Error in cancel callback for %s: %s", run_id, e)
return True
async def register_run(self, run_id: str, name: Optional[str] = None, cancel_callback: Optional[CancelCallback] = None) -> None:
"""Register a run for future cancellation.
Args:
run_id: Unique run identifier (string)
name: Optional friendly name for debugging/observability
cancel_callback: Optional async callback called when a cancel is requested
"""
async with self._lock:
self._runs[run_id] = {"name": name, "registered_at": time.time(), "cancel_callback": cancel_callback, "cancelled": False}
logger.info("Registered run %s (%s)", run_id, name)
async def unregister_run(self, run_id: str) -> None:
"""Remove a run from tracking.
Args:
run_id: Unique identifier for the run to unregister.
"""
async with self._lock:
if run_id in self._runs:
self._runs.pop(run_id, None)
logger.info("Unregistered run %s", run_id)
async def cancel_run(self, run_id: str, reason: Optional[str] = None) -> bool:
"""Attempt to cancel a run.
Args:
run_id: Unique identifier for the run to cancel.
reason: Optional textual reason for the cancellation request.
Returns:
bool: True if the run was found and cancellation was attempted (or already marked),
False if the run was not known locally.
"""
cancel_cb = None
entry = None
async with self._lock:
entry = self._runs.get(run_id)
if not entry:
# Entry not found - will publish to Redis outside the lock
pass
elif entry.get("cancelled"):
logger.debug("Run %s already cancelled", run_id)
return True
else:
entry["cancelled"] = True
entry["cancelled_at"] = time.time()
entry["cancel_reason"] = reason
cancel_cb = entry.get("cancel_callback")
# Handle unknown run case outside the lock
if not entry:
logger.info("Cancellation requested for unknown run %s (queued for remote peers)", run_id)
# Publish to Redis for other workers (outside lock to avoid blocking)
await self._publish_cancellation(run_id, reason)
return False
# Log cancellation with reason and request_id for observability
logger.info("Tool execution cancelled: run_id=%s, reason=%s, tool=%s", run_id, reason or "not specified", entry.get("name", "unknown"))
if cancel_cb:
try:
await cancel_cb(reason)
logger.info("Cancel callback executed for %s", run_id)
except Exception as e:
logger.exception("Error in cancel callback for %s: %s", run_id, e)
# Publish to Redis for other workers
await self._publish_cancellation(run_id, reason)
return True
async def _publish_cancellation(self, run_id: str, reason: Optional[str] = None) -> None:
"""Publish cancellation event to Redis for other workers.
Args:
run_id: Unique identifier for the run being cancelled.
reason: Optional textual reason for the cancellation.
"""
if not self._redis:
return
try:
message = json.dumps({"run_id": run_id, "reason": reason})
await self._redis.publish("cancellation:cancel", message)
logger.debug("Published cancellation to Redis: run_id=%s", run_id)
except Exception as e:
logger.warning(f"Failed to publish cancellation to Redis: {e}")
async def get_status(self, run_id: str) -> Optional[Dict[str, Any]]:
"""Return the status dict for a run if known, else None.
Args:
run_id: Unique identifier for the run to query.
Returns:
Optional[Dict[str, Any]]: The status dictionary for the run if found, otherwise None.
"""
async with self._lock:
return self._runs.get(run_id)
async def is_registered(self, run_id: str) -> bool:
"""Check if a run is currently registered.
Args:
run_id: Unique identifier for the run to check.
Returns:
bool: True if the run is registered, False otherwise.
"""
async with self._lock:
return run_id in self._runs
# Module-level singleton for importers to use
cancellation_service = CancellationService()
