Keep task routing keys alive via sliding TTL; version-aware worker server resolution

A resumed leg that runs longer than its pointer's wall-clock TTL stranded
_lookup_task on the base leg (false completion / not found). Each poll now
refreshes the routing keys' TTL (sliding expiration), so an actively-polled
task keeps them alive regardless of execution duration, and the resumed-leg
write uses the same buffered TTL as creation. Separately, remote-worker server
resolution now respects the requested tool version, so two versions of the same
mounted tool name resolve to their own child server.

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):

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)

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).

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

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):

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