diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index 22bdf0259..8a9c421b7 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -145,15 +145,22 @@ class TaskContextSnapshot: origin_request_id: str | None = None session_id: str | None = None owning_tool_name: str | None = None + owning_tool_version: str | None = None @classmethod - def capture(cls, owning_tool_name: str | None = None) -> TaskContextSnapshot: + def capture( + cls, + owning_tool_name: str | None = None, + owning_tool_version: str | None = None, + ) -> TaskContextSnapshot: """Capture current context for background task execution. - ``owning_tool_name`` is the routable name of the tool the call targeted. - A remote worker (separate process) cannot reach the submitting process's - server map, so it re-resolves the owning (child) server from this name - against the root — see ``make_task_context``. + ``owning_tool_name``/``owning_tool_version`` identify the exact tool the + call targeted. A remote worker (separate process) cannot reach the + submitting process's server map, so it re-resolves the owning (child) + server from this name and version against the root — see + ``make_task_context``. The version matters when two versions of the same + mounted tool name live on different child servers. """ from fastmcp.server.dependencies import ( get_access_token, @@ -178,6 +185,7 @@ class TaskContextSnapshot: ), session_id=session_id, owning_tool_name=owning_tool_name, + owning_tool_version=owning_tool_version, ) @classmethod @@ -195,6 +203,7 @@ class TaskContextSnapshot: origin_request_id=parsed.get("origin_request_id"), session_id=parsed.get("session_id"), owning_tool_name=parsed.get("owning_tool_name"), + owning_tool_version=parsed.get("owning_tool_version"), ) def to_json(self) -> str: @@ -206,6 +215,7 @@ class TaskContextSnapshot: "origin_request_id": self.origin_request_id, "session_id": self.session_id, "owning_tool_name": self.owning_tool_name, + "owning_tool_version": self.owning_tool_version, } ) @@ -435,10 +445,19 @@ async def _resolve_owning_server( from fastmcp.exceptions import NotFoundError from fastmcp.server.dependencies import get_server from fastmcp.server.providers.fastmcp_provider import FastMCPProviderTool + from fastmcp.utilities.versions import VersionSpec root = get_server() + # Resolve the exact version the call targeted: two versions of the same + # mounted tool name can live on different child servers, so omitting the + # version could pick the wrong server's state and masking policy. + version = ( + VersionSpec(eq=snapshot.owning_tool_version) + if snapshot.owning_tool_version + else None + ) try: - tool = await root.get_tool(snapshot.owning_tool_name) + tool = await root.get_tool(snapshot.owning_tool_name, version) except NotFoundError: return None if isinstance(tool, FastMCPProviderTool): diff --git a/fastmcp_tasks/fastmcp_tasks/creation.py b/fastmcp_tasks/fastmcp_tasks/creation.py index 745e7ceb7..7fc65ead9 100644 --- a/fastmcp_tasks/fastmcp_tasks/creation.py +++ b/fastmcp_tasks/fastmcp_tasks/creation.py @@ -109,7 +109,9 @@ async def create_task( created_at_key = docket.key(f"{prefix}:{task_id}:created_at") poll_interval_key = docket.key(f"{prefix}:{task_id}:poll_interval") - snapshot = TaskContextSnapshot.capture(owning_tool_name=tool.name) + snapshot = TaskContextSnapshot.capture( + owning_tool_name=tool.name, owning_tool_version=tool.version + ) async with docket.redis() as redis: await redis.set(task_meta_key, task_key, ex=ttl_seconds) diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py index a4f7b0800..b12aabf19 100644 --- a/fastmcp_tasks/fastmcp_tasks/handlers.py +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -33,7 +33,11 @@ from fastmcp.tools.base import InputRequiredToolResult, Tool, ToolResult from fastmcp.utilities.tasks import DEFAULT_POLL_INTERVAL_MS from fastmcp.utilities.versions import VersionSpec from fastmcp_tasks.context import get_task_scope -from fastmcp_tasks.creation import enqueue_task_leg, registered_component_for_key +from fastmcp_tasks.creation import ( + TASK_MAPPING_TTL_BUFFER_SECONDS, + enqueue_task_leg, + registered_component_for_key, +) from fastmcp_tasks.input_store import ( acquire_update_lock, acquire_update_lock_blocking, @@ -43,6 +47,7 @@ from fastmcp_tasks.input_store import ( load_task_args, mark_cancelled, read_outstanding_inputs, + refresh_current_leg_ttl, release_update_lock, save_current_leg, store_input_responses, @@ -108,6 +113,16 @@ def _ttl_ms(docket: Docket) -> int: return int(docket.execution_ttl.total_seconds() * 1000) +def _task_key_ttl_seconds(docket: Docket) -> int: + """Wall-clock TTL for a task's Redis metadata keys. + + Docket's ``execution_ttl`` plus a buffer (matching task creation), so a key + written or refreshed now comfortably outlives the execution-retention + window. Sliding expiration on each poll keeps it alive for long legs. + """ + return int(docket.execution_ttl.total_seconds()) + TASK_MAPPING_TTL_BUFFER_SECONDS + + async def _lookup_task( docket: Docket, task_scope: str | None, task_id: str ) -> tuple[Any, str, int, str | None, int]: @@ -141,6 +156,16 @@ async def _lookup_task( if not execution: raise _task_not_found(task_id) + # Sliding expiration: an actively-polled task refreshes its routing keys so + # they never expire mid-execution — a resumed leg that runs longer than the + # keys' wall-clock TTL would otherwise strand `_lookup_task` on the base leg. + refresh_ttl = _task_key_ttl_seconds(docket) + async with docket.redis() as redis: + await redis.expire(meta_key, refresh_ttl) + await redis.expire(created_at_key, refresh_ttl) + await redis.expire(poll_key, refresh_ttl) + await refresh_current_leg_ttl(docket, task_scope, task_id, refresh_ttl) + created_at = created_at_bytes.decode("utf-8") if created_at_bytes else None try: @@ -350,9 +375,13 @@ async def tasks_update( next_leg_key = leg_execution_key(base_task_key, next_leg) await enqueue_task_leg(server, docket, component, raw_arguments, next_leg_key) - ttl_seconds = int(docket.execution_ttl.total_seconds()) await save_current_leg( - docket, task_scope, task_id, next_leg_key, next_leg, ttl_seconds + docket, + task_scope, + task_id, + next_leg_key, + next_leg, + _task_key_ttl_seconds(docket), ) # The answered leg's surfaced keys are now superseded; drop them so they # are never reused (SEP-2663 L350). diff --git a/fastmcp_tasks/fastmcp_tasks/input_store.py b/fastmcp_tasks/fastmcp_tasks/input_store.py index edd46ef64..d13a6e82e 100644 --- a/fastmcp_tasks/fastmcp_tasks/input_store.py +++ b/fastmcp_tasks/fastmcp_tasks/input_store.py @@ -184,6 +184,22 @@ async def save_current_leg( ) +async def refresh_current_leg_ttl( + docket: Docket, task_scope: str | None, task_id: str, ttl_seconds: int +) -> None: + """Extend the current-leg pointer's TTL (sliding expiration). + + The pointer is written with a wall-clock TTL, but a leg's execution can run + longer than that — a resumed guard leg especially. Refreshing on each poll + keeps the routing pointer alive for an actively-polled task no matter how + long the leg runs, so ``_lookup_task`` never falls back to the base leg + while the current leg is still executing. + """ + async with docket.redis() as redis: + await redis.expire(_current_leg_key(docket, task_scope, task_id), ttl_seconds) + await redis.expire(_leg_number_key(docket, task_scope, task_id), ttl_seconds) + + async def load_current_leg( docket: Docket, task_scope: str | None, task_id: str ) -> tuple[str | None, int]: diff --git a/tests/tasks/server/test_task_mount.py b/tests/tasks/server/test_task_mount.py index c13836247..df8f23deb 100644 --- a/tests/tasks/server/test_task_mount.py +++ b/tests/tasks/server/test_task_mount.py @@ -167,6 +167,47 @@ class TestRemoteWorkerServerResolution: finally: _current_server.reset(token) + async def test_resolve_owning_server_respects_version(self): + """Two versions of a mounted tool name resolve to their own child server.""" + import weakref + + from fastmcp_tasks.context import ( + TaskContextSnapshot, + _resolve_owning_server, + ) + + from fastmcp.server.dependencies import _current_server + + child_v1 = FastMCP("child-v1") + + @child_v1.tool(name="calc", version="1.0", task=True) + async def calc_v1() -> str: + return "v1" + + child_v2 = FastMCP("child-v2") + + @child_v2.tool(name="calc", version="2.0", task=True) + async def calc_v2() -> str: + return "v2" + + parent = FastMCP("parent-versions") + parent.add_extension(TasksExtension()) + parent.mount(child_v1) + parent.mount(child_v2) + + token = _current_server.set(weakref.ref(parent)) + try: + resolved_v1 = await _resolve_owning_server( + TaskContextSnapshot(owning_tool_name="calc", owning_tool_version="1.0") + ) + resolved_v2 = await _resolve_owning_server( + TaskContextSnapshot(owning_tool_name="calc", owning_tool_version="2.0") + ) + assert resolved_v1 is child_v1 + assert resolved_v2 is child_v2 + finally: + _current_server.reset(token) + class TestMountedToolTasksNoPrefix: async def test_mounted_tool_without_prefix_works(self, child_server): diff --git a/tests/tasks/server/test_task_ttl.py b/tests/tasks/server/test_task_ttl.py index a3bb74b2d..8fa9f670e 100644 --- a/tests/tasks/server/test_task_ttl.py +++ b/tests/tasks/server/test_task_ttl.py @@ -69,3 +69,30 @@ async def test_default_ttl_when_unspecified(): assert created.ttl_ms == DEFAULT_TTL_MS got = await get_task(mcp, created.task_id) assert got.ttl_ms == DEFAULT_TTL_MS + + +async def test_poll_refreshes_routing_key_ttl(): + """A poll extends the current-leg pointer's TTL (sliding expiration). + + A leg that runs longer than the pointer's wall-clock TTL would otherwise + strand `_lookup_task` on the base leg. Polling must keep the routing keys + alive: after shrinking the pointer's TTL, a `tasks/get` restores it. + """ + from fastmcp_tasks.input_store import _current_leg_key + + mcp = _ttl_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "slow_task", {}) + docket = mcp._docket + assert docket is not None + key = _current_leg_key(docket, None, created.task_id) + + async with docket.redis() as redis: + await redis.expire(key, 5) + assert await redis.ttl(key) <= 5 + + await get_task(mcp, created.task_id) + + async with docket.redis() as redis: + # Refreshed well past the shrunk 5s, back toward the full window. + assert await redis.ttl(key) > 60