diff --git a/CLAUDE.md b/CLAUDE.md index 70fb9b1..fc3297d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -134,8 +134,8 @@ The server requires Stacklet credentials configured through one of: **Server Settings:** Additional configuration via environment variables with `STACKLET_MCP_` prefix: -- `STACKLET_MCP_DOWNLOADS_PATH` (default: system temp directory) - Directory for storing query result files -- `STACKLET_MCP_DOWNLOADS_ENABLED` (default: true) - Write query result files; set false for hosted deployments +- `STACKLET_MCP_DOWNLOADS_PATH` (default: system temp directory) - Directory for storing query result and dataset export files +- `STACKLET_MCP_DOWNLOADS_ENABLED` (default: true) - Write query result and dataset export files; set false for hosted deployments - `STACKLET_MCP_ASSETDB_DATASOURCE_NAME` (default: AssetDB) - Name the data source id is looked up from on first use - `STACKLET_MCP_ASSETDB_DATASOURCE` (default: unset) - AssetDB data source ID, when a deployment needs to name one explicitly; unset means look it up by name - `STACKLET_MCP_ASSETDB_ALLOW_SAVE` (default: false) - Enable query save/update functionality @@ -153,6 +153,11 @@ available by exactly one of two routes, depending on `downloads_enabled`: attached to the tool response as a resource. Hosted deployments must disable downloads — a caller elsewhere cannot read the server's filesystem, so the files would only accumulate. +Dataset export tools follow the same switch. Enabled, a completed export is downloaded with +the caller's token to `export_*` in the downloads directory (`full_results_saved_to`). +Disabled, the tool returns the link with a `download_note` to open it in a signed-in browser. +An expired export returns only a `download_note` either way. + **External Dependencies:** - Documentation files are fetched from the live Stacklet docs service at runtime - Redash endpoint is derived by replacing "api." with "redash." in the platform endpoint diff --git a/NEWS.md b/NEWS.md index 5922090..be520ea 100644 --- a/NEWS.md +++ b/NEWS.md @@ -4,6 +4,12 @@ ### Changes +- **Dataset exports download with your credentials**: Stacklet now requires sign-in to + download an export, and only the user who started it can. When the server writes files, + `platform_dataset_lookup` and `platform_dataset_export` download the finished export with + your token and return the local path in `full_results_saved_to`. A hosted server returns + the link with a `download_note` telling you to open it in your signed-in browser. + ### Fixes --- diff --git a/README.md b/README.md index 62528a8..babfc0b 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ When the MCP is run from an agent, those can be set in the `"env"` section of th The following variables are available: -- `STACKLET_MCP_DOWNLOADS_PATH`: directory for storing query result files (default: system temp directory) +- `STACKLET_MCP_DOWNLOADS_PATH`: directory for storing query result and dataset export files (default: system temp directory) - `STACKLET_MCP_ASSETDB_DATASOURCE_NAME`: the name of the AssetDB data source in Redash, which the id is looked up from on first use (default: `AssetDB`) - `STACKLET_MCP_ASSETDB_DATASOURCE`: the datasource ID for AssetDB in Redash, for a deployment that needs to name one explicitly. Unset by default, meaning look it up by name — Redash assigns the id at creation, so it differs between deployments and cannot be assumed - `STACKLET_MCP_ASSETDB_ALLOW_SAVE`: whether to enable write operations in AssetDB (default: `false`) diff --git a/stacklet/mcp/platform/dataset_info.md b/stacklet/mcp/platform/dataset_info.md index 12f2d37..7f0539e 100644 --- a/stacklet/mcp/platform/dataset_info.md +++ b/stacklet/mcp/platform/dataset_info.md @@ -85,7 +85,10 @@ platform_dataset_export( **Key status fields:** `dataset_id`, `success`, `download_url`, `available_until` (24hr expiry) -**Download:** `curl -L -o file.csv "download_url"` (follows S3 redirects) +**Download:** the link needs the credentials of the user who started the export, so do not fetch it yourself. +- `full_results_saved_to` set: the server already downloaded the file to that local path. +- `download_note` set and `available_until` still ahead: give the user `download_url` and the note. They open it in a browser where they're signed in to Stacklet. +- `available_until` passed: the link is dead. Start a new export instead of sharing it. ## Quick Examples @@ -110,6 +113,6 @@ platform_dataset_export( - Cannot use paging params (`first`, `after`) - export handles pagination - `timeout=0` returns immediately; use `platform_dataset_lookup` to monitor -- Download URLs redirect to S3 - use `curl -L` +- Only the user who started an export can download it, while signed in - Files expire after 24 hours - Use GraphQL Node IDs for node-based exports diff --git a/stacklet/mcp/platform/graphql.py b/stacklet/mcp/platform/graphql.py index c052450..54e89b3 100644 --- a/stacklet/mcp/platform/graphql.py +++ b/stacklet/mcp/platform/graphql.py @@ -8,6 +8,7 @@ """ import asyncio +import os import re import time @@ -33,6 +34,7 @@ from ..settings import SETTINGS from ..stacklet_auth import StackletCredentials from ..utils.error import AnnotatedError +from ..utils.file import download_file from .models import ( ConnectionExport, ExportRequest, @@ -225,8 +227,37 @@ async def _get_export(self, dataset_id: str) -> ConnectionExport: # If no errors, data is at least guaranteed guaranteed truthy. fields = cast(dict[str, Any], result.data)["node"] + if fields is None: + raise AnnotatedError( + problem=f"Dataset export {dataset_id} not found", + likely_cause="the ID is wrong, or the export belongs to another user", + next_steps="start a new export with platform_dataset_export", + ) return ConnectionExport(**fields) + async def download_export(self, dataset_id: str, url: str) -> str: + """Stream an export file into the downloads directory, returning its path.""" + try: + async with self.session.stream("GET", url) as response: + if not response.is_success: + await response.aread() + response.raise_for_status() + return await _save_stream(response) + except httpx.HTTPError as e: + detail = f"{e}\n{e.response.text}" if isinstance(e, httpx.HTTPStatusError) else str(e) + raise AnnotatedError( + problem=f"Downloading dataset export {dataset_id} failed", + likely_cause=( + "the export expired, these credentials cannot read it, " + "or the connection dropped" + ), + next_steps=( + f"retry with platform_dataset_lookup, or give the user {url} " + "to open in a browser where they're signed in to Stacklet" + ), + original_error=detail, + ) from e + Q_GET_EXPORT = """ query getExport($id: ID!) { node(id: $id) { @@ -267,6 +298,19 @@ async def _query(self, query: str, variables: dict[str, Any]) -> GraphQLQueryRes raise Exception(f"Unexpected response: {response.text}") +async def _save_stream(response: httpx.Response) -> str: + """Write a response body to a new downloads CSV, removing it if anything fails.""" + f = download_file("wb", "export_", ".csv") + try: + with f: + async for chunk in response.aiter_bytes(): + await asyncio.to_thread(f.write, chunk) + except BaseException: + os.unlink(f.name) + raise + return f.name + + def has_mutations(doc: DocumentNode) -> bool: """Return whether a GraphQL document calls mutations.""" operations = {dd.operation for dd in doc.definitions} diff --git a/stacklet/mcp/platform/models.py b/stacklet/mcp/platform/models.py index 00ee1f9..bb5cf92 100644 --- a/stacklet/mcp/platform/models.py +++ b/stacklet/mcp/platform/models.py @@ -150,3 +150,15 @@ class ConnectionExport(BaseModel): available_until: datetime | None = Field( None, validation_alias="availableUntil", description="When download URL expires" ) + + # A succeeded export with a download URL sets exactly one of these. + full_results_saved_to: str | None = Field( + None, + description=( + "Local path where the server saved the exported file. Null when the server " + "does not write files or the export has expired, in which case see download_note." + ), + ) + download_note: str | None = Field( + None, description="Why the server did not download the file, and what to do instead" + ) diff --git a/stacklet/mcp/platform/tools.py b/stacklet/mcp/platform/tools.py index b1bf6a0..ad40abd 100644 --- a/stacklet/mcp/platform/tools.py +++ b/stacklet/mcp/platform/tools.py @@ -3,6 +3,7 @@ # Copyright (c) 2025-2026 Stacklet, Inc. # +from datetime import UTC, datetime from typing import Annotated, Any from fastmcp import Context @@ -25,6 +26,26 @@ ) +HOSTED_DOWNLOAD_NOTE = ( + "Open this link in a browser where you're signed in to Stacklet. " + "It works only for you, until the export expires." +) +EXPIRED_DOWNLOAD_NOTE = "This export has expired. Start a new one with platform_dataset_export." + + +async def _deliver(client: PlatformClient, export: ConnectionExport) -> ConnectionExport: + """Download a completed export's file, or say how to use its link instead.""" + if not export.completed or not export.download_url: + return export + # The platform keeps returning the URL after the file has expired. + if export.available_until and export.available_until <= datetime.now(UTC): + return export.model_copy(update={"download_note": EXPIRED_DOWNLOAD_NOTE}) + if not SETTINGS.downloads_enabled: + return export.model_copy(update={"download_note": HOSTED_DOWNLOAD_NOTE}) + path = await client.download_export(export.dataset_id, export.download_url) + return export.model_copy(update={"full_results_saved_to": path}) + + def tools() -> list[Tool]: """List of available Platform tools.""" return [ @@ -37,7 +58,7 @@ def tools() -> list[Tool]: # Exports don't change governance data, but they do start a job and # write a file server-side, and each call starts another one. make_tool(platform_dataset_export, read_only=False, destructive=False, idempotent=False), - make_tool(platform_dataset_lookup, read_only=True), + make_tool(platform_dataset_lookup, read_only=False, destructive=False, idempotent=False), ] @@ -211,7 +232,10 @@ async def platform_dataset_export( 1. Define columns mapping GraphQL fields to CSV columns 2. Optionally add filters via params 3. Export runs asynchronously - use timeout=0 to return immediately - 4. Use platform_dataset_lookup() to check progress and get download URL + 4. Use platform_dataset_lookup() to check progress and get the result + + A completed export's file is downloaded before the call returns when the server writes + files, so a large export takes longer than the timeout alone suggests. """ dataset_input = ExportRequest( connection_field=connection_field, @@ -222,7 +246,7 @@ async def platform_dataset_export( client = PlatformClient.get(ctx) dataset_id = await client.start_export(dataset_input) - return await client.wait_for_export(dataset_id, timeout) + return await _deliver(client, await client.wait_for_export(dataset_id, timeout)) @json_guard @@ -251,11 +275,14 @@ async def platform_dataset_lookup( Export states: - Processing: Export is running (shows progress if available) - - Complete: Ready for download (includes download_url and expiry time) + - Complete: Ready for download (includes download_url and expiry time, plus either + full_results_saved_to or download_note; an expired export has only the note) - Failed: Export encountered an error - Set timeout > 0 to wait for completion, or timeout=0 for immediate status check. - Download URLs are temporary and expire after a few hours. + Set timeout > 0 to wait for completion, or timeout=0 for an immediate status check. + Every lookup of a completed export downloads its file again when the server writes + files, so the call takes as long as that download. + Download URLs expire after 24 hours and work only for the user who started the export. """ client = PlatformClient.get(ctx) - return await client.wait_for_export(dataset_id, timeout) + return await _deliver(client, await client.wait_for_export(dataset_id, timeout)) diff --git a/stacklet/mcp/settings.py b/stacklet/mcp/settings.py index 0d7d162..93d2764 100644 --- a/stacklet/mcp/settings.py +++ b/stacklet/mcp/settings.py @@ -23,10 +23,10 @@ class Settings(BaseSettings): downloads_enabled: bool = Field( default=True, description=( - "Write complete query results to downloads_path. Disable for hosted " - "deployments, where the caller cannot read the server's filesystem and " - "the files would accumulate unbounded; full results are then attached " - "to the tool response instead." + "Write complete query results and dataset exports to downloads_path. " + "Disable for hosted deployments, where the caller cannot read the server's " + "filesystem and the files would accumulate unbounded; query results are then " + "attached to the tool response, and exports return their link with a note." ), ) diff --git a/tests/test_tools_platform.py b/tests/test_tools_platform.py index 31fc9a9..f33a87d 100644 --- a/tests/test_tools_platform.py +++ b/tests/test_tools_platform.py @@ -7,6 +7,8 @@ Tests for Platform MCP tools using FastMCP's in-memory testing pattern. """ +import json + from unittest.mock import ANY, MagicMock import httpx @@ -17,9 +19,10 @@ from stacklet.mcp.platform.graphql import PlatformClient, has_mutations from stacklet.mcp.platform.models import ExportParam -from stacklet.mcp.platform.tools import tools +from stacklet.mcp.platform.tools import EXPIRED_DOWNLOAD_NOTE, HOSTED_DOWNLOAD_NOTE, tools +from stacklet.mcp.settings import SETTINGS -from .testing.http import ExpectRequest +from .testing.http import ExpectRequest, MockHTTPXResponse from .testing.mcp import MCPBearerTest, MCPTest, json_guard_parametrize @@ -344,11 +347,16 @@ def graphql_field_error(message: str, field_path: list, line: int = 1, column: i } +AVAILABLE_UNTIL = "2099-12-07T03:15:09+00:00" +EXPIRED = "2024-12-07T03:15:09+00:00" + + class PlatformDatasetTest(MCPBearerTest): DATASET_ID = "node-123" + CSV = "id\nabc\n" @staticmethod - def dataset_result(dataset_id, started=False, succeeded=None): + def dataset_result(dataset_id, started=False, succeeded=None, available_until=AVAILABLE_UNTIL): node = { "id": dataset_id, "started": None, @@ -368,7 +376,7 @@ def dataset_result(dataset_id, started=False, succeeded=None): node |= { "message": "yay!", "downloadURL": "https://example.com/x.csv", - "availableUntil": "2024-12-07T03:15:09+00:00", + "availableUntil": available_until, } else: node |= {"message": "meh."} @@ -407,7 +415,9 @@ def fancy_export_param(): }, } - def assert_result(self, result, started=False, succeeded=None): + def assert_result( + self, result, started=False, succeeded=None, expired=False, downloads_enabled=True + ): """Assert that result matches expected dataset state from factory args.""" expected = { "dataset_id": self.DATASET_ID, @@ -429,11 +439,42 @@ def assert_result(self, result, started=False, succeeded=None): if succeeded: expected["message"] = "yay!" expected["download_url"] = "https://example.com/x.csv" - expected["available_until"] = "2024-12-07T03:15:09Z" + expected["available_until"] = (EXPIRED if expired else AVAILABLE_UNTIL).replace( + "+00:00", "Z" + ) else: expected["message"] = "meh." - assert result.json() == expected + actual = result.json() + saved_to = actual.pop("full_results_saved_to") + note = actual.pop("download_note") + assert actual == expected + + # A delivered export reaches the caller by exactly one route. + delivered = succeeded and not expired + if delivered and downloads_enabled: + assert note is None + with open(saved_to) as f: + assert f.read() == self.CSV + else: + assert saved_to is None + if delivered: + assert note == HOSTED_DOWNLOAD_NOTE + else: + assert note == (EXPIRED_DOWNLOAD_NOTE if expired else None) + assert list(SETTINGS.downloads_path.iterdir()) == [] + + def expect_download(self, status_code=200, response=None): + """Create expectation for the export file download, made with the bearer token.""" + return ExpectRequest( + "https://example.com/x.csv", + status_code=status_code, + response=self.CSV if response is None else response, + ) + + def delivery(self, succeeded): + """Expected requests that download a completed export's file.""" + return [self.expect_download()] if succeeded else [] def expect_start_export(self, columns, connection="someConnection", node_id=None, params=None): """Create expectation for the export mutation request.""" @@ -550,6 +591,7 @@ async def test_response_complete(self): with self.http.expect( self.expect_start_export(columns), self.expect_get_export(dataset), + *self.delivery(succeeded=True), ): result = await self.assert_call( { @@ -638,6 +680,7 @@ async def test_complete_timeout_immediate(self, succeeded, mangle, value, async_ with self.http.expect( self.expect_get_export(complete), + *self.delivery(succeeded), ): result = await self.assert_call( {"dataset_id": self.DATASET_ID, "timeout": mangle(value)} @@ -657,6 +700,7 @@ async def test_complete_timeout_delayed(self, succeeded, mangle, value, async_sl self.expect_get_export(incomplete), self.expect_get_export(incomplete), self.expect_get_export(complete), + *self.delivery(succeeded), ): result = await self.assert_call( {"dataset_id": self.DATASET_ID, "timeout": mangle(value)} @@ -665,6 +709,79 @@ async def test_complete_timeout_delayed(self, succeeded, mangle, value, async_sl assert async_sleeps == [2, 4] self.assert_result(result, started=True, succeeded=succeeded) + async def test_not_found(self): + # The platform returns a null node for an unknown ID, and for an export that + # belongs to someone else. + not_found = self.expect_get_export({"id": self.DATASET_ID}) + not_found.response = json.dumps({"data": {"node": None}}) + + with self.http.expect(not_found): + result = await self.assert_call({"dataset_id": self.DATASET_ID}, error=True) + + assert result.text.startswith(f"Dataset export {self.DATASET_ID} not found") + + @pytest.mark.parametrize("status_code", [401, 404]) + async def test_download_refused(self, status_code): + complete = self.dataset_result(self.DATASET_ID, started=True, succeeded=True) + body = json.dumps({"error": "Export not found or expired."}) + + with self.http.expect( + self.expect_get_export(complete), + self.expect_download(status_code=status_code, response=body), + ): + result = await self.assert_call({"dataset_id": self.DATASET_ID}, error=True) + + assert f"Downloading dataset export {self.DATASET_ID} failed" in result.text + assert str(status_code) in result.text + assert "Export not found or expired." in result.text + assert "https://example.com/x.csv" in result.text + assert list(SETTINGS.downloads_path.iterdir()) == [] + + async def test_download_interrupted(self, monkeypatch): + async def broken_stream(self): + yield b"id\n" + raise httpx.ReadError("connection lost") + + monkeypatch.setattr(MockHTTPXResponse, "aiter_bytes", broken_stream) + complete = self.dataset_result(self.DATASET_ID, started=True, succeeded=True) + + with self.http.expect(self.expect_get_export(complete), self.expect_download()): + result = await self.assert_call({"dataset_id": self.DATASET_ID}, error=True) + + # The caller keeps what it needs to retry: the export ID and the link. + assert f"Downloading dataset export {self.DATASET_ID} failed" in result.text + assert "https://example.com/x.csv" in result.text + assert list(SETTINGS.downloads_path.iterdir()) == [] + + @pytest.mark.parametrize( + "downloads_enabled, succeeded, expired", + [ + pytest.param(False, True, False, id="hosted"), + pytest.param(False, False, False, id="hosted-failed"), + pytest.param(False, True, True, id="hosted-expired"), + pytest.param(True, True, True, id="local-expired"), + ], + ) + async def test_no_download(self, override_setting, downloads_enabled, succeeded, expired): + override_setting("downloads_enabled", downloads_enabled) + complete = self.dataset_result( + self.DATASET_ID, + started=True, + succeeded=succeeded, + available_until=EXPIRED if expired else AVAILABLE_UNTIL, + ) + + with self.http.expect(self.expect_get_export(complete)): + result = await self.assert_call({"dataset_id": self.DATASET_ID}) + + self.assert_result( + result, + started=True, + succeeded=succeeded, + expired=expired, + downloads_enabled=downloads_enabled, + ) + class TestToolAnnotations: """The behavioural hints hosts show when asking to approve a tool call.""" @@ -681,7 +798,6 @@ def annotations(self, name: str) -> ToolAnnotations: ("platform_dataset_info", False), ("platform_graphql_list_types", True), ("platform_graphql_get_types", True), - ("platform_dataset_lookup", True), ], ) def test_read_only(self, name: str, open_world: bool): @@ -690,8 +806,10 @@ def test_read_only(self, name: str, open_world: bool): assert annotations.destructiveHint is False assert annotations.openWorldHint is open_world - def test_dataset_export_writes_but_destroys_nothing(self): - annotations = self.annotations("platform_dataset_export") + @pytest.mark.parametrize("name", ["platform_dataset_export", "platform_dataset_lookup"]) + def test_dataset_tools_write_but_destroy_nothing(self, name: str): + # Both download a completed export into the downloads directory. + annotations = self.annotations(name) assert annotations.readOnlyHint is False assert annotations.destructiveHint is False assert annotations.idempotentHint is False diff --git a/tests/testing/http.py b/tests/testing/http.py index 1138837..a7c9833 100644 --- a/tests/testing/http.py +++ b/tests/testing/http.py @@ -7,6 +7,7 @@ import json +from contextlib import asynccontextmanager from typing import Any import httpx @@ -28,9 +29,19 @@ def content(self): def text(self): return self._data + @property + def is_success(self): + return 200 <= self.status_code < 300 + def json(self): return json.loads(self._data) + async def aread(self): + return self.content + + async def aiter_bytes(self): + yield self.content + def raise_for_status(self): if self.status_code >= 400: raise httpx.HTTPStatusError( @@ -112,7 +123,13 @@ async def mock_request(self, method, url, **kwargs): auth_check_func(self, mock_stacklet_credentials) return controller.next_request().respond(method, url, **kwargs) + @asynccontextmanager + async def mock_stream(self, method, url, **kwargs): + auth_check_func(self, mock_stacklet_credentials) + yield controller.next_request().respond(method, url, **kwargs) + monkeypatch.setattr("httpx.AsyncClient.request", mock_request) + monkeypatch.setattr("httpx.AsyncClient.stream", mock_stream) return controller