mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 13:04:18 +02:00
Refactor task module: rename protocol.py to requests.py and reduce redundancy (#2667)
Renamed protocol.py to requests.py for clarity - it handles MCP task request endpoints (tasks/get, tasks/result, tasks/cancel, tasks/list). Consolidated constants to config.py as single source of truth: - DEFAULT_POLL_INTERVAL_MS now derived from DEFAULT_POLL_INTERVAL - TaskConfig uses constant instead of hardcoded timedelta(seconds=5) Extracted _lookup_task_execution() helper to eliminate ~50 lines of duplicated Redis lookup code. Uses redis.mget() for single round-trip instead of 3 separate calls (performance improvement).
This commit is contained in:
parent
7ff23a06f4
commit
3afdb3f7b3
4 changed files with 90 additions and 86 deletions
|
|
@ -606,7 +606,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
ServerResult,
|
||||
)
|
||||
|
||||
from fastmcp.server.tasks.protocol import (
|
||||
from fastmcp.server.tasks.requests import (
|
||||
tasks_cancel_handler,
|
||||
tasks_get_handler,
|
||||
tasks_list_handler,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,11 @@ from typing import Any, Literal
|
|||
# Task execution modes per SEP-1686 / MCP ToolExecution.taskSupport
|
||||
TaskMode = Literal["forbidden", "optional", "required"]
|
||||
|
||||
# Default values for task metadata (single source of truth)
|
||||
DEFAULT_POLL_INTERVAL = timedelta(seconds=5) # Default poll interval
|
||||
DEFAULT_POLL_INTERVAL_MS = int(DEFAULT_POLL_INTERVAL.total_seconds() * 1000)
|
||||
DEFAULT_TTL_MS = 60_000 # Default TTL in milliseconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskConfig:
|
||||
|
|
@ -47,7 +52,7 @@ class TaskConfig:
|
|||
"""
|
||||
|
||||
mode: TaskMode = "optional"
|
||||
poll_interval: timedelta = timedelta(seconds=5)
|
||||
poll_interval: timedelta = DEFAULT_POLL_INTERVAL
|
||||
|
||||
@classmethod
|
||||
def from_bool(cls, value: bool) -> TaskConfig:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""SEP-1686 task protocol handlers.
|
||||
"""SEP-1686 task request handlers.
|
||||
|
||||
Implements MCP task protocol methods: tasks/get, tasks/result, tasks/list, tasks/cancel, tasks/delete.
|
||||
Handles MCP task protocol requests: tasks/get, tasks/result, tasks/list, tasks/cancel.
|
||||
These handlers query and manage existing tasks (contrast with handlers.py which creates tasks).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -20,6 +21,7 @@ from mcp.types import (
|
|||
ListTasksResult,
|
||||
)
|
||||
|
||||
from fastmcp.server.tasks.config import DEFAULT_POLL_INTERVAL_MS, DEFAULT_TTL_MS
|
||||
from fastmcp.server.tasks.keys import parse_task_key
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -37,6 +39,69 @@ DOCKET_TO_MCP_STATE: dict[ExecutionState, str] = {
|
|||
}
|
||||
|
||||
|
||||
async def _lookup_task_execution(
|
||||
docket: Any,
|
||||
session_id: str,
|
||||
client_task_id: str,
|
||||
) -> tuple[Any, str | None, int]:
|
||||
"""Look up task execution and metadata from Redis.
|
||||
|
||||
Consolidates the common pattern of fetching task metadata from Redis,
|
||||
validating it exists, and retrieving the Docket execution.
|
||||
|
||||
Args:
|
||||
docket: Docket instance
|
||||
session_id: Session ID
|
||||
client_task_id: Client-provided task ID
|
||||
|
||||
Returns:
|
||||
Tuple of (execution, created_at, poll_interval_ms)
|
||||
|
||||
Raises:
|
||||
McpError: If task not found or execution not found
|
||||
"""
|
||||
# Build Redis keys
|
||||
redis_key = f"fastmcp:task:{session_id}:{client_task_id}"
|
||||
created_at_key = f"{redis_key}:created_at"
|
||||
poll_interval_key = f"{redis_key}:poll_interval"
|
||||
|
||||
# Fetch metadata (single round-trip with mget)
|
||||
async with docket.redis() as redis:
|
||||
task_key_bytes, created_at_bytes, poll_interval_bytes = await redis.mget(
|
||||
redis_key, created_at_key, poll_interval_key
|
||||
)
|
||||
|
||||
# Decode and validate task_key
|
||||
task_key = task_key_bytes.decode("utf-8") if task_key_bytes else None
|
||||
if not task_key:
|
||||
raise McpError(
|
||||
ErrorData(code=INVALID_PARAMS, message=f"Task {client_task_id} not found")
|
||||
)
|
||||
|
||||
# Get execution
|
||||
execution = await docket.get_execution(task_key)
|
||||
if not execution:
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message=f"Task {client_task_id} execution not found",
|
||||
)
|
||||
)
|
||||
|
||||
# Parse metadata with defaults
|
||||
created_at = created_at_bytes.decode("utf-8") if created_at_bytes else None
|
||||
try:
|
||||
poll_interval_ms = (
|
||||
int(poll_interval_bytes.decode("utf-8"))
|
||||
if poll_interval_bytes
|
||||
else DEFAULT_POLL_INTERVAL_MS
|
||||
)
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
poll_interval_ms = DEFAULT_POLL_INTERVAL_MS
|
||||
|
||||
return execution, created_at, poll_interval_ms
|
||||
|
||||
|
||||
async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskResult:
|
||||
"""Handle MCP 'tasks/get' request (SEP-1686).
|
||||
|
||||
|
|
@ -61,7 +126,7 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
|
|||
# Get session ID from Context
|
||||
session_id = ctx.session_id
|
||||
|
||||
# Get execution from Docket (use instance attribute for cross-task access)
|
||||
# Get Docket instance
|
||||
docket = server._docket
|
||||
if docket is None:
|
||||
raise McpError(
|
||||
|
|
@ -71,45 +136,10 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
|
|||
)
|
||||
)
|
||||
|
||||
# Look up task metadata from Redis
|
||||
redis_key = f"fastmcp:task:{session_id}:{client_task_id}"
|
||||
created_at_key = f"fastmcp:task:{session_id}:{client_task_id}:created_at"
|
||||
poll_interval_key = f"fastmcp:task:{session_id}:{client_task_id}:poll_interval"
|
||||
async with docket.redis() as redis:
|
||||
task_key_bytes = await redis.get(redis_key)
|
||||
created_at_bytes = await redis.get(created_at_key)
|
||||
poll_interval_bytes = await redis.get(poll_interval_key)
|
||||
|
||||
task_key = None if task_key_bytes is None else task_key_bytes.decode("utf-8")
|
||||
created_at = (
|
||||
None if created_at_bytes is None else created_at_bytes.decode("utf-8")
|
||||
# Look up task execution and metadata
|
||||
execution, created_at, poll_interval_ms = await _lookup_task_execution(
|
||||
docket, session_id, client_task_id
|
||||
)
|
||||
try:
|
||||
poll_interval_ms = (
|
||||
int(poll_interval_bytes.decode("utf-8"))
|
||||
if poll_interval_bytes
|
||||
else 5000 # Default to 5 seconds
|
||||
)
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
poll_interval_ms = 5000
|
||||
|
||||
if task_key is None:
|
||||
# Task not found - raise error per MCP protocol
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS, message=f"Task {client_task_id} not found"
|
||||
)
|
||||
)
|
||||
|
||||
execution = await docket.get_execution(task_key)
|
||||
if execution is None:
|
||||
# Task key exists but no execution - raise error
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message=f"Task {client_task_id} execution not found",
|
||||
)
|
||||
)
|
||||
|
||||
# Sync state from Redis
|
||||
await execution.sync()
|
||||
|
|
@ -138,7 +168,7 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
|
|||
status=mcp_state, # type: ignore[arg-type]
|
||||
createdAt=created_at, # type: ignore[arg-type]
|
||||
lastUpdatedAt=datetime.now(timezone.utc),
|
||||
ttl=60000,
|
||||
ttl=DEFAULT_TTL_MS,
|
||||
pollInterval=poll_interval_ms,
|
||||
statusMessage=status_message,
|
||||
)
|
||||
|
|
@ -345,6 +375,7 @@ async def tasks_cancel_handler(
|
|||
# Get session ID from Context
|
||||
session_id = ctx.session_id
|
||||
|
||||
# Get Docket instance
|
||||
docket = server._docket
|
||||
if docket is None:
|
||||
raise McpError(
|
||||
|
|
@ -354,48 +385,14 @@ async def tasks_cancel_handler(
|
|||
)
|
||||
)
|
||||
|
||||
# Look up task metadata from Redis
|
||||
redis_key = f"fastmcp:task:{session_id}:{client_task_id}"
|
||||
created_at_key = f"fastmcp:task:{session_id}:{client_task_id}:created_at"
|
||||
poll_interval_key = f"fastmcp:task:{session_id}:{client_task_id}:poll_interval"
|
||||
async with docket.redis() as redis:
|
||||
task_key_bytes = await redis.get(redis_key)
|
||||
created_at_bytes = await redis.get(created_at_key)
|
||||
poll_interval_bytes = await redis.get(poll_interval_key)
|
||||
|
||||
task_key = None if task_key_bytes is None else task_key_bytes.decode("utf-8")
|
||||
created_at = (
|
||||
None if created_at_bytes is None else created_at_bytes.decode("utf-8")
|
||||
# Look up task execution and metadata
|
||||
execution, created_at, poll_interval_ms = await _lookup_task_execution(
|
||||
docket, session_id, client_task_id
|
||||
)
|
||||
try:
|
||||
poll_interval_ms = (
|
||||
int(poll_interval_bytes.decode("utf-8"))
|
||||
if poll_interval_bytes
|
||||
else 5000 # Default to 5 seconds
|
||||
)
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
poll_interval_ms = 5000
|
||||
|
||||
if task_key is None:
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message=f"Invalid taskId: {client_task_id} not found",
|
||||
)
|
||||
)
|
||||
|
||||
# Check if task exists
|
||||
execution = await docket.get_execution(task_key)
|
||||
if execution is None:
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message=f"Invalid taskId: {client_task_id} not found",
|
||||
)
|
||||
)
|
||||
|
||||
# Cancel via Docket (now sets CANCELLED state natively)
|
||||
await docket.cancel(task_key)
|
||||
# Note: We need to get task_key from execution.key for cancellation
|
||||
await docket.cancel(execution.key)
|
||||
|
||||
# Return task status with cancelled state
|
||||
# createdAt is REQUIRED per SEP-1686 final spec (line 430)
|
||||
|
|
@ -403,9 +400,11 @@ async def tasks_cancel_handler(
|
|||
return CancelTaskResult(
|
||||
taskId=client_task_id,
|
||||
status="cancelled",
|
||||
createdAt=created_at or datetime.now(timezone.utc).isoformat(),
|
||||
createdAt=datetime.fromisoformat(created_at)
|
||||
if created_at
|
||||
else datetime.now(timezone.utc),
|
||||
lastUpdatedAt=datetime.now(timezone.utc),
|
||||
ttl=60_000,
|
||||
ttl=DEFAULT_TTL_MS,
|
||||
pollInterval=poll_interval_ms,
|
||||
statusMessage="Task cancelled",
|
||||
)
|
||||
|
|
@ -13,7 +13,7 @@ from typing import TYPE_CHECKING
|
|||
from docket.execution import ExecutionState
|
||||
from mcp.types import TaskStatusNotification, TaskStatusNotificationParams
|
||||
|
||||
from fastmcp.server.tasks.protocol import DOCKET_TO_MCP_STATE
|
||||
from fastmcp.server.tasks.requests import DOCKET_TO_MCP_STATE
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue