46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
"""
|
|
Shared polling utilities for task operations.
|
|
|
|
This module provides generic polling logic that works for both client→server
|
|
and server→client task polling.
|
|
|
|
WARNING: These APIs are experimental and may change without notice.
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
|
|
|
import anyio
|
|
|
|
from mcp.shared.experimental.tasks.helpers import is_terminal
|
|
from mcp.types import GetTaskResult
|
|
|
|
|
|
async def poll_until_terminal(
|
|
get_task: Callable[[str], Awaitable[GetTaskResult]],
|
|
task_id: str,
|
|
default_interval_ms: int = 500,
|
|
) -> AsyncIterator[GetTaskResult]:
|
|
"""
|
|
Poll a task until it reaches terminal status.
|
|
|
|
This is a generic utility that works for both client→server and server→client
|
|
polling. The caller provides the get_task function appropriate for their direction.
|
|
|
|
Args:
|
|
get_task: Async function that takes task_id and returns GetTaskResult
|
|
task_id: The task to poll
|
|
default_interval_ms: Fallback poll interval if server doesn't specify
|
|
|
|
Yields:
|
|
GetTaskResult for each poll
|
|
"""
|
|
while True:
|
|
status = await get_task(task_id)
|
|
yield status
|
|
|
|
if is_terminal(status.status):
|
|
break
|
|
|
|
interval_ms = status.pollInterval if status.pollInterval is not None else default_interval_ms
|
|
await anyio.sleep(interval_ms / 1000)
|