Skip to content

Commit 6b9dbda

Browse files
RafaelPogithub-actions[bot]
RafaelPo
authored and
github-actions[bot]
committed
fix(mcp): route download through public API instead of internal endpoint (#5001)
## Summary - MCP download route was calling the Engine's internal `/tasks/{id}/output_rows` by stripping `/api/v0` from the public URL — this goes through Traefik which doesn't expose internal routes, causing 404 - Fix: forward the per-task API key from Redis and call the public `GET /api/v0/tasks/{id}/result` endpoint instead - Stop popping the task token on completion so it remains available for downloads (Redis TTL expires it naturally) Fixes https://futuresearch-ai.sentry.io/issues/7357940590/ ## Test plan - [x] All 19 route tests pass - [ ] Deploy MCP to staging, verify CSV download works end-to-end 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Sourced from commit 77721e7adc1ca5227769aacef8ab6e7c55dd4f93
1 parent 90f2494 commit 6b9dbda

4 files changed

Lines changed: 71 additions & 61 deletions

File tree

‎futuresearch-mcp/src/futuresearch_mcp/routes.py‎

Lines changed: 22 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import json
77
import logging
88
import secrets
9-
from typing import Any
109
from uuid import UUID
1110

1211
import httpx
@@ -173,8 +172,8 @@ async def api_progress(request: Request) -> Response: # noqa: PLR0911
173172

174173
ts = TaskState(status_response)
175174

176-
if ts.is_terminal:
177-
await redis_store.pop_task_token(task_id)
175+
# Don't pop the token on completion — the download route needs it.
176+
# Let the Redis TTL expire it naturally.
178177

179178
return JSONResponse(
180179
ts.model_dump(mode="json", exclude=_UI_EXCLUDE), headers=cors
@@ -236,15 +235,14 @@ async def api_download_url(request: Request) -> Response:
236235
return JSONResponse({"download_url": download_url}, headers=cors)
237236

238237

239-
async def api_download(request: Request) -> Response:
238+
async def api_download(request: Request) -> Response: # noqa: PLR0911
240239
"""Download task results as CSV or JSON.
241240
242-
Unauthenticated — the task ID (UUID) is sufficient. The Engine's
243-
internal ``/tasks/{id}/output_rows`` endpoint is also unauthenticated.
241+
Fetches results from the public Engine API using the per-task API key
242+
stored in Redis.
244243
245244
Query params:
246245
format: "csv" (default) or "json"
247-
raw: "true" to keep _source_bank for client-side citation rendering
248246
"""
249247
cors = _cors_headers()
250248
if request.method == "OPTIONS":
@@ -263,30 +261,28 @@ async def api_download(request: Request) -> Response:
263261
return JSONResponse(
264262
{"error": "Unsupported format"}, status_code=400, headers=cors
265263
)
266-
raw = request.query_params.get("raw", "").lower() in ("true", "1")
267-
268-
# Pass processing flags to the Engine — it does the stripping/resolution.
269-
engine_params: dict[str, Any] = {"offset": 0, "limit": 100000}
270-
if raw:
271-
engine_params["strip_internal_cols"] = True
272-
else:
273-
engine_params.update(
274-
resolve_citations=True,
275-
strip_source_bank=True,
276-
strip_internal_cols=True,
264+
# Fetch results via the public API (paginated path handles citation
265+
# resolution and internal column stripping automatically).
266+
api_key = await redis_store.get_task_token(task_id)
267+
if not api_key:
268+
return JSONResponse(
269+
{"error": "Unknown task or expired session"},
270+
status_code=404,
271+
headers=cors,
277272
)
278273

279274
try:
280-
engine_base = settings.futuresearch_api_url.removesuffix(
281-
"/api/v0"
282-
).removesuffix("/")
283-
async with httpx.AsyncClient(base_url=engine_base) as http:
284-
resp = await http.post(
285-
f"/tasks/{task_id}/output_rows",
286-
params=engine_params,
275+
async with httpx.AsyncClient(
276+
base_url=settings.futuresearch_api_url,
277+
headers={"Authorization": f"Bearer {api_key}"},
278+
) as http:
279+
resp = await http.get(
280+
f"/tasks/{task_id}/result",
281+
params={"offset": 0, "limit": 100000},
287282
)
288283
resp.raise_for_status()
289-
records: list[dict] = resp.json()
284+
body = resp.json()
285+
records: list[dict] = body.get("data") or []
290286
except Exception:
291287
logger.exception("Failed to fetch results for download, task %s", task_id)
292288
return JSONResponse(

‎futuresearch-mcp/tests/test_http_integration.py‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,8 @@ async def test_completed_task_cleans_up_tokens(self, client: httpx.AsyncClient):
225225
assert resp.status_code == 200
226226
assert resp.json()["status"] == "completed"
227227

228-
# Task token cleaned up; poll token kept for CSV download
229-
assert await redis_store.get_task_token(task_id) is None
228+
# Both tokens kept — task token needed for CSV download, TTL expires them
229+
assert await redis_store.get_task_token(task_id) is not None
230230
assert await redis_store.get_poll_token(task_id) is not None
231231

232232
@pytest.mark.asyncio
@@ -299,12 +299,8 @@ async def test_progress_lifecycle(self, client: httpx.AsyncClient):
299299
assert resp.status_code == 200
300300
assert resp.json()["status"] == "completed"
301301

302-
# 4. Task token is gone — further progress polls return 404
303-
resp = await client.get(
304-
f"/api/progress/{task_id}",
305-
params={"token": poll_token},
306-
)
307-
assert resp.status_code == 404
302+
# 4. Task token is still available (TTL-based expiry, not popped)
303+
assert await redis_store.get_task_token(task_id) is not None
308304

309305

310306
# ── Download-token endpoint ──────────────────────────────────
@@ -387,14 +383,18 @@ async def test_multiple_calls_return_same_url(self, client: httpx.AsyncClient):
387383

388384
@pytest.mark.asyncio
389385
async def test_download_is_repeatable(self, client: httpx.AsyncClient):
390-
"""Download endpoint can be called multiple times (unauthenticated)."""
386+
"""Download endpoint can be called multiple times."""
391387
task_id = str(uuid4())
388+
await redis_store.store_task_token(task_id, "sk-cho-test")
392389

393390
mock_resp = MagicMock()
394391
mock_resp.raise_for_status = MagicMock()
395-
mock_resp.json.return_value = [{"x": 1, "y": 2}]
392+
mock_resp.json.return_value = {
393+
"data": [{"x": 1, "y": 2}],
394+
"status": "completed",
395+
}
396396
mock_http = AsyncMock()
397-
mock_http.__aenter__.return_value.post.return_value = mock_resp
397+
mock_http.__aenter__.return_value.get.return_value = mock_resp
398398

399399
with patch(
400400
"futuresearch_mcp.routes.httpx.AsyncClient",
@@ -423,6 +423,7 @@ async def test_get_url_and_download(self, client: httpx.AsyncClient):
423423

424424
await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
425425
await redis_store.store_task_owner(task_id, "test-user")
426+
await redis_store.store_task_token(task_id, "sk-cho-test")
426427

427428
# 1. Get the download URL
428429
mint_resp = await client.get(
@@ -433,15 +434,18 @@ async def test_get_url_and_download(self, client: httpx.AsyncClient):
433434
download_url = mint_resp.json()["download_url"]
434435
assert f"/api/results/{task_id}/download" in download_url
435436

436-
# 2. Download CSV — Engine is called to fetch rows (unauthenticated)
437+
# 2. Download CSV via public API (forwards API key from Redis)
437438
mock_resp = MagicMock()
438439
mock_resp.raise_for_status = MagicMock()
439-
mock_resp.json.return_value = [
440-
{"name": "Alice", "score": 95},
441-
{"name": "Bob", "score": 87},
442-
]
440+
mock_resp.json.return_value = {
441+
"data": [
442+
{"name": "Alice", "score": 95},
443+
{"name": "Bob", "score": 87},
444+
],
445+
"status": "completed",
446+
}
443447
mock_http = AsyncMock()
444-
mock_http.__aenter__.return_value.post.return_value = mock_resp
448+
mock_http.__aenter__.return_value.get.return_value = mock_resp
445449

446450
with patch(
447451
"futuresearch_mcp.routes.httpx.AsyncClient",

‎futuresearch-mcp/tests/test_result_store.py‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -290,15 +290,19 @@ async def client(self, app):
290290
@pytest.mark.asyncio
291291
async def test_valid_download(self, client: httpx.AsyncClient):
292292
task_id = str(uuid4())
293+
await redis_store.store_task_token(task_id, "sk-cho-test")
293294

294295
mock_resp = MagicMock()
295296
mock_resp.raise_for_status = MagicMock()
296-
mock_resp.json.return_value = [
297-
{"name": "Alice", "score": 95},
298-
{"name": "Bob", "score": 87},
299-
]
297+
mock_resp.json.return_value = {
298+
"data": [
299+
{"name": "Alice", "score": 95},
300+
{"name": "Bob", "score": 87},
301+
],
302+
"status": "completed",
303+
}
300304
mock_http = AsyncMock()
301-
mock_http.__aenter__.return_value.post.return_value = mock_resp
305+
mock_http.__aenter__.return_value.get.return_value = mock_resp
302306

303307
with patch(
304308
"futuresearch_mcp.routes.httpx.AsyncClient",
@@ -315,15 +319,19 @@ async def test_valid_download(self, client: httpx.AsyncClient):
315319
async def test_json_format_returns_json(self, client: httpx.AsyncClient):
316320
"""?format=json returns a JSON array fetched from the Engine."""
317321
task_id = str(uuid4())
322+
await redis_store.store_task_token(task_id, "sk-cho-test")
318323

319324
mock_resp = MagicMock()
320325
mock_resp.raise_for_status = MagicMock()
321-
mock_resp.json.return_value = [
322-
{"name": "Alice", "score": 95},
323-
{"name": "Bob", "score": 87},
324-
]
326+
mock_resp.json.return_value = {
327+
"data": [
328+
{"name": "Alice", "score": 95},
329+
{"name": "Bob", "score": 87},
330+
],
331+
"status": "completed",
332+
}
325333
mock_http = AsyncMock()
326-
mock_http.__aenter__.return_value.post.return_value = mock_resp
334+
mock_http.__aenter__.return_value.get.return_value = mock_resp
327335

328336
with patch(
329337
"futuresearch_mcp.routes.httpx.AsyncClient",

‎futuresearch-mcp/tests/test_routes.py‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -281,8 +281,8 @@ async def test_completed_task_pops_tokens(self):
281281
body = json.loads(bytes(resp.body).decode())
282282
assert body["status"] == "completed"
283283

284-
# Task token cleaned up; poll token kept for CSV download
285-
assert await redis_store.get_task_token(task_id) is None
284+
# Both tokens kept — task token needed for CSV download, TTL expires them
285+
assert await redis_store.get_task_token(task_id) is not None
286286
assert await redis_store.get_poll_token(task_id) is not None
287287

288288
@pytest.mark.asyncio
@@ -380,6 +380,7 @@ async def test_download_url_works_for_download(self):
380380

381381
await redis_store.store_poll_token(task_id, poll_token, user_id="test-user")
382382
await redis_store.store_task_owner(task_id, "test-user")
383+
await redis_store.store_task_token(task_id, "sk-cho-test")
383384

384385
# Step 1: Get the download URL
385386
mint_req = FakeRequest(
@@ -392,15 +393,15 @@ async def test_download_url_works_for_download(self):
392393
download_url = mint_body["download_url"]
393394
assert f"/api/results/{task_id}/download" in download_url
394395

395-
# Step 2: Download directly (unauthenticated)
396+
# Step 2: Download via public API (forwards API key from Redis)
396397
dl_req = FakeRequest(
397398
path_params={"task_id": task_id},
398399
)
399400
mock_resp = MagicMock()
400401
mock_resp.raise_for_status = MagicMock()
401-
mock_resp.json.return_value = records
402+
mock_resp.json.return_value = {"data": records, "status": "completed"}
402403
mock_http = AsyncMock()
403-
mock_http.__aenter__.return_value.post.return_value = mock_resp
404+
mock_http.__aenter__.return_value.get.return_value = mock_resp
404405

405406
with patch(
406407
"futuresearch_mcp.routes.httpx.AsyncClient",
@@ -427,15 +428,16 @@ async def test_csv_download_quotes_all_fields(self):
427428
},
428429
]
429430

430-
# Download CSV directly (unauthenticated)
431+
await redis_store.store_task_token(task_id, "sk-cho-test")
432+
431433
dl_req = FakeRequest(
432434
path_params={"task_id": task_id},
433435
)
434436
mock_resp = MagicMock()
435437
mock_resp.raise_for_status = MagicMock()
436-
mock_resp.json.return_value = records
438+
mock_resp.json.return_value = {"data": records, "status": "completed"}
437439
mock_http = AsyncMock()
438-
mock_http.__aenter__.return_value.post.return_value = mock_resp
440+
mock_http.__aenter__.return_value.get.return_value = mock_resp
439441

440442
with patch(
441443
"futuresearch_mcp.routes.httpx.AsyncClient",

0 commit comments

Comments
 (0)