refactor(mcp): remove redundant Redis ownership checks (#4994)

## Summary

- Remove `_check_task_ownership()` from the MCP tool layer (tools.py) —
the Engine already validates task ownership on every API call via
`get_user_task() → session.owner_account_id == user_id`
- Remove `_validate_task_owner()` cross-check from REST endpoints
(routes.py) — the poll token is cryptographically random and
task-scoped, making the ownership cross-check redundant defense-in-depth
- Remove 4 Redis functions (`store_task_owner`, `get_task_owner`,
`get_poll_token_owner`, `user_id` param from `store_poll_token`) and
eliminate 2 Redis keys per task: `mcp:task_owner:{task_id}` and
`mcp:poll_owner:{task_id}`

Net: **-319 lines**, fewer Redis round-trips per request, no change in
security posture.

## Test plan

- [x] All 371 existing tests pass (format, lint, typecheck hooks all
green)
- [x] Ownership-specific tests removed (tested behavior no longer
exists)
- [x] Poll token validation remains intact as primary REST auth gate
- [ ] Verify on staging: submit task → poll progress → fetch results
works end-to-end
- [ ] Verify multi-user isolation still works via Engine auth

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Sourced from commit c5fc538f4b3e8c68d4c3268e4417c9711ff334f9

Author: RafaelPo

futuresearch-mcp/src/futuresearch_mcp/redis_store.py

@@ -192,42 +192,17 @@ async def pop_task_token(task_id: str) -> None:
 # ── Poll tokens ───────────────────────────────────────────────
 
 
-async def store_poll_token(task_id: str, poll_token: str, user_id: str = "") -> None:
-    """Store an encrypted poll token, optionally bound to a user identity."""
-    client = get_redis_client()
-    await client.setex(
+async def store_poll_token(task_id: str, poll_token: str) -> None:
+    """Store an encrypted poll token."""
+    await get_redis_client().setex(
         name=build_key("poll_token", task_id),
         time=TOKEN_TTL,
         value=encrypt_value(poll_token),
     )
-    if user_id:
-        await client.setex(
-            name=build_key("poll_owner", task_id),
-            time=TOKEN_TTL,
-            value=user_id,
-        )
 
 
 async def get_poll_token(task_id: str) -> str | None:
     encrypted = await get_redis_client().get(build_key("poll_token", task_id))
     if encrypted is None:
         return None
     return decrypt_value(encrypted)
-
-
-async def get_poll_token_owner(task_id: str) -> str | None:
-    """Return the user_id bound to a poll token, or None if not set."""
-    return await get_redis_client().get(build_key("poll_owner", task_id))
-
-
-# ── Task ownership (user-scoped data isolation) ──────────────
-
-
-async def store_task_owner(task_id: str, user_id: str) -> None:
-    """Record which user submitted a task (for cross-user access checks)."""
-    await get_redis_client().setex(build_key("task_owner", task_id), TOKEN_TTL, user_id)
-
-
-async def get_task_owner(task_id: str) -> str | None:
-    """Return the user_id that owns a task, or None if not recorded."""
-    return await get_redis_client().get(build_key("task_owner", task_id))

futuresearch-mcp/src/futuresearch_mcp/routes.py

@@ -81,57 +81,7 @@ async def _validate_poll_token(task_id: str, request: Request) -> JSONResponse |
     return None
 
 
-async def _validate_task_owner(task_id: str) -> JSONResponse | None:
-    """Verify the task has a recorded owner and that the poll token was
-    issued for the same user.  Returns a 403 response if ownership cannot
-    be verified, or ``None`` if the caller may proceed.
-
-    Fail-closed: tasks without an ownership record are rejected.  When a
-    poll-token owner is recorded, it must match the task owner — this
-    cross-check detects ownership-record tampering and ensures the poll
-    token was legitimately issued for this task/user pair.
-    """
-    owner = await redis_store.get_task_owner(task_id)
-    if not owner:
-        logger.warning(
-            "REST access denied for task %s: no owner recorded (fail-closed)",
-            task_id,
-        )
-        return JSONResponse(
-            {"error": "Task ownership could not be verified"},
-            status_code=403,
-            headers=_cors_headers(),
-        )
-
-    poll_owner = await redis_store.get_poll_token_owner(task_id)
-    if not poll_owner:
-        logger.warning(
-            "REST access denied for task %s: no poll_owner recorded (fail-closed)",
-            task_id,
-        )
-        return JSONResponse(
-            {"error": "Task ownership could not be verified"},
-            status_code=403,
-            headers=_cors_headers(),
-        )
-    if poll_owner != owner:
-        logger.warning(
-            "REST access denied for task %s: poll_owner=%s != task_owner=%s",
-            task_id,
-            poll_owner,
-            owner,
-        )
-        return JSONResponse(
-            {"error": "Task ownership could not be verified"},
-            status_code=403,
-            headers=_cors_headers(),
-        )
-
-    logger.debug("REST access granted for task %s (owner=%s)", task_id, owner)
-    return None
-
-
-async def api_progress(request: Request) -> Response:  # noqa: PLR0911
+async def api_progress(request: Request) -> Response:
     """REST endpoint for the session widget to poll task progress."""
     cors = _cors_headers()
     if request.method == "OPTIONS":
@@ -148,9 +98,6 @@ async def api_progress(request: Request) -> Response:  # noqa: PLR0911
     if err := await _validate_poll_token(task_id, request):
         return err
 
-    if err := await _validate_task_owner(task_id):
-        return err
-
     api_key = await redis_store.get_task_token(task_id)
 
     if not api_key:
@@ -228,9 +175,6 @@ async def api_download_url(request: Request) -> Response:
     if err := await _validate_poll_token_bearer_only(task_id, request):
         return err
 
-    if err := await _validate_task_owner(task_id):
-        return err
-
     download_url = f"{settings.mcp_server_url}/api/results/{task_id}/download"
     return JSONResponse({"download_url": download_url}, headers=cors)
 

futuresearch-mcp/src/futuresearch_mcp/tool_helpers.py

@@ -20,7 +20,6 @@
 from futuresearch.generated.models.task_status import TaskStatus
 from futuresearch.generated.models.task_status_response import TaskStatusResponse
 from futuresearch.generated.types import Unset
-from mcp.server.auth.middleware.auth_context import get_access_token
 from mcp.server.fastmcp import Context
 from mcp.server.session import ServerSession
 from mcp.types import TextContent
@@ -199,33 +198,16 @@ def _submission_text(label: str, task_id: str, session_id: str = "") -> str:
 
 
 async def _record_task_ownership(task_id: str, token: str) -> str:
-    """Record task ownership and create a poll token.
+    """Store the API token and create a poll token for a submitted task.
 
     Must run for every HTTP submission (including internal clients) so that
-    downstream ownership checks in progress/results don't fail.
+    downstream poll-token checks in progress/results don't fail.
 
     Returns the poll_token.
     """
     poll_token = secrets.token_urlsafe(32)
     await redis_store.store_task_token(task_id, token)
-
-    # Record task owner for cross-user access checks (HTTP mode only).
-    # This MUST succeed — downstream ownership checks deny access when no
-    # owner is recorded, so a silent failure here would lock the user out
-    # of their own task.
-    user_id = ""
-    if settings.is_http:
-        access_token = get_access_token()
-        if not access_token or not access_token.client_id:
-            raise RuntimeError(
-                f"Cannot record task owner for {task_id}: no authenticated user"
-            )
-        user_id = access_token.client_id
-        await redis_store.store_task_owner(task_id, user_id)
-
-    # Bind the poll token to the same user identity so the REST layer
-    # can cross-check poll_owner == task_owner.
-    await redis_store.store_poll_token(task_id, poll_token, user_id=user_id)
+    await redis_store.store_poll_token(task_id, poll_token)
     return poll_token
 
 

futuresearch-mcp/src/futuresearch_mcp/tools.py

@@ -27,7 +27,6 @@
 )
 from futuresearch.session import create_session, list_sessions
 from futuresearch.task import cancel_task
-from mcp.server.auth.middleware.auth_context import get_access_token
 from mcp.types import CallToolResult, TextContent, ToolAnnotations
 from pydantic import BaseModel, create_model
 
@@ -83,50 +82,6 @@ def _error_result(text: str) -> CallToolResult:
     )
 
 
-async def _check_task_ownership(task_id: str) -> list[TextContent] | None:
-    """Verify the current user owns *task_id*. Returns an error response if
-    access should be denied, or ``None`` if the caller may proceed.
-
-    Only active in HTTP mode; always returns ``None`` for stdio.
-
-    When no owner is recorded in Redis (e.g. tasks created via the presigned
-    upload URL flow, which bypasses the MCP tool layer), the current user is
-    auto-registered as owner.  The Engine independently validates ownership
-    on every API call via session-level checks, so this is safe — a user
-    cannot claim a task they don't own because subsequent Engine calls would
-    fail.
-    """
-    if not settings.is_http:
-        return None
-
-    access_token = get_access_token()
-    user_id = access_token.client_id if access_token else None
-    if not user_id:
-        return [TextContent(type="text", text="Access denied: no authenticated user.")]
-
-    owner = await redis_store.get_task_owner(task_id)
-    if not owner:
-        # Task was likely created outside the MCP tool layer (e.g. presigned
-        # URL upload).  Claim it for the current user — the Engine will
-        # independently reject any API calls if this user doesn't actually
-        # own the task's session.
-        logger.info(
-            "No owner recorded for task %s — auto-registering user %s",
-            task_id,
-            user_id,
-        )
-        await redis_store.store_task_owner(task_id, user_id)
-        return None
-
-    if user_id != owner:
-        return [
-            TextContent(
-                type="text", text="Access denied: this task belongs to another user."
-            )
-        ]
-    return None
-
-
 @mcp.tool(
     name="futuresearch_browse_lists",
     structured_output=False,
@@ -958,18 +913,6 @@ async def futuresearch_progress(
     client = _get_client(ctx)
     task_id = params.task_id
 
-    # ── Cross-user access check ──────────────────────────────────
-    try:
-        if denied := await _check_task_ownership(task_id):
-            return denied
-    except Exception:
-        logger.exception("Could not verify task ownership for %s", task_id)
-        return [
-            TextContent(
-                type="text", text="Unable to verify task ownership. Please try again."
-            )
-        ]
-
     # Block server-side before polling — controls the cadence
     await asyncio.sleep(redis_store.PROGRESS_POLL_DELAY)
 
@@ -1110,14 +1053,6 @@ async def futuresearch_results_http(
     if is_internal_client():
         skip_widget = True
 
-    # ── Cross-user access check ──────────────────────────────────
-    try:
-        if denied := await _check_task_ownership(task_id):
-            return CallToolResult(content=denied, isError=True)  # pyright: ignore[reportArgumentType]  # list invariance
-    except Exception:
-        logger.exception("Could not verify task ownership for %s", task_id)
-        return _error_result("Unable to verify task ownership. Please try again.")
-
     # ── Fetch paginated rows directly from Engine ─────────────────
     try:
         rows, total_count, session_id, artifact_id = await _fetch_task_result(
@@ -1369,18 +1304,6 @@ async def futuresearch_cancel(
     client = _get_client(ctx)
     task_id = params.task_id
 
-    # ── Cross-user access check ──────────────────────────────────
-    try:
-        if denied := await _check_task_ownership(task_id):
-            return denied
-    except Exception:
-        logger.exception("Could not verify task ownership for %s", task_id)
-        return [
-            TextContent(
-                type="text", text="Unable to verify task ownership. Please try again."
-            )
-        ]
-
     try:
         await cancel_task(task_id=UUID(task_id), client=client)
         return [

futuresearch-mcp/tests/test_http_integration.py

@@ -142,25 +142,11 @@ async def test_unauthorized_with_wrong_token(self, client: httpx.AsyncClient):
         assert resp.status_code == 403
 
     @pytest.mark.asyncio
-    async def test_denied_without_owner(self, client: httpx.AsyncClient):
-        """Poll token is valid but no task owner recorded → fail-closed 403."""
+    async def test_unknown_task_returns_404(self, client: httpx.AsyncClient):
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
         await redis_store.store_poll_token(task_id, poll_token)
-        # No task owner stored — ownership check should reject
 
-        resp = await client.get(
-            f"/api/progress/{task_id}", params={"token": poll_token}
-        )
-        assert resp.status_code == 403
-        assert resp.json()["error"] == "Task ownership could not be verified"
-
-    @pytest.mark.asyncio
-    async def test_unknown_task_returns_404(self, client: httpx.AsyncClient):
-        task_id = str(uuid4())
-        poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
-        await redis_store.store_task_owner(task_id, "test-user")
         # No task_token stored
 
         resp = await client.get(
@@ -173,9 +159,8 @@ async def test_unknown_task_returns_404(self, client: httpx.AsyncClient):
     async def test_running_task_returns_progress(self, client: httpx.AsyncClient):
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
         await redis_store.store_task_token(task_id, "api-key-123")
-        await redis_store.store_task_owner(task_id, "test-user")
 
         status_resp = _make_status_response(
             status="running", completed=7, total=20, failed=1, running=4
@@ -205,9 +190,8 @@ async def test_running_task_returns_progress(self, client: httpx.AsyncClient):
     async def test_completed_task_cleans_up_tokens(self, client: httpx.AsyncClient):
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
         await redis_store.store_task_token(task_id, "api-key")
-        await redis_store.store_task_owner(task_id, "test-user")
 
         status_resp = _make_status_response(
             status="completed", completed=10, total=10, failed=0, running=0
@@ -233,9 +217,8 @@ async def test_completed_task_cleans_up_tokens(self, client: httpx.AsyncClient):
     async def test_api_error_returns_500(self, client: httpx.AsyncClient):
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
         await redis_store.store_task_token(task_id, "api-key")
-        await redis_store.store_task_owner(task_id, "test-user")
 
         with patch(
             "futuresearch_mcp.routes.get_task_status_tasks_task_id_status_get.asyncio",
@@ -264,9 +247,8 @@ async def test_progress_lifecycle(self, client: httpx.AsyncClient):
         poll_token = secrets.token_urlsafe(16)
 
         # 1. Store tokens (simulating what futuresearch_agent does)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
         await redis_store.store_task_token(task_id, "api-key-for-polling")
-        await redis_store.store_task_owner(task_id, "test-user")
 
         # 2. Poll progress — task is running
         running_resp = _make_status_response(status="running", completed=1, total=3)
@@ -318,8 +300,7 @@ async def test_bearer_auth_via_real_http_header(self, client: httpx.AsyncClient)
         """Authorization: Bearer header is parsed correctly by Starlette."""
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
-        await redis_store.store_task_owner(task_id, "test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
 
         resp = await client.get(
             f"/api/results/{task_id}/download-token",
@@ -350,8 +331,7 @@ async def test_query_param_rejected_through_real_url(
         """Poll token via ?token= query param in a real URL is rejected."""
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
-        await redis_store.store_task_owner(task_id, "test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
 
         resp = await client.get(
             f"/api/results/{task_id}/download-token",
@@ -364,8 +344,7 @@ async def test_multiple_calls_return_same_url(self, client: httpx.AsyncClient):
         """Two sequential calls for the same task return the same download URL."""
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
-        await redis_store.store_task_owner(task_id, "test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
 
         headers = {"Authorization": f"Bearer {poll_token}"}
         resp1 = await client.get(
@@ -421,8 +400,7 @@ async def test_get_url_and_download(self, client: httpx.AsyncClient):
         task_id = str(uuid4())
         poll_token = secrets.token_urlsafe(16)
 
-        await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
-        await redis_store.store_task_owner(task_id, "test-user")
+        await redis_store.store_poll_token(task_id, poll_token)
         await redis_store.store_task_token(task_id, "sk-cho-test")
 
         # 1. Get the download URL