Session-aware HPC plot wrappers (#25)

CHANGELOG.md

@@ -50,6 +50,11 @@ under their entry.
   endpoint state (no endpoint configured, or endpoint not ready) instead
   of surfacing a misleading `FileNotFoundError` from the local fallback.
   Local paths and HEALPix specs still fall back silently. (#27)
+- `plot_mesh`, `plot_variable`, and `plot_zonal_mean` HPC wrappers now
+  accept `session_id` + `dataset_handle` and resolve grid/data paths
+  from a registered session dataset, matching the handle pattern
+  already used by `subset_bbox` / `subset_polygon`. Direct paths still
+  work unchanged. (#25)
 
 ## 0.1.0 — initial scaffold
 

src/uxarray_mcp/tools/remote_tools.py

@@ -377,12 +377,13 @@ def calculate_zonal_mean_hpc(
 
 
 def plot_mesh_hpc(
-    grid_path: str,
+    grid_path: str | None = None,
     width: int = 800,
     height: int = 400,
     use_remote: bool = False,
     endpoint: str | None = None,
     session_id: str | None = None,
+    dataset_handle: str | None = None,
 ) -> list[Any]:
     """Render a mesh wireframe PNG with optional HPC execution.
 
@@ -392,15 +393,19 @@ def plot_mesh_hpc(
 
     Parameters
     ----------
-    grid_path : str
+    grid_path : str | None
         Path to mesh file. Can be a local path or an HPC filesystem path
-        when use_remote=True.
+        when use_remote=True. Optional when ``session_id`` and
+        ``dataset_handle`` are provided.
     width : int
         Image width in pixels (default 800).
     height : int
         Image height in pixels (default 400).
     use_remote : bool
         If True and HPC is configured, render on the remote endpoint.
+    session_id, dataset_handle : str | None
+        When both are provided, the grid path is looked up from the
+        registered session dataset instead of ``grid_path``.
 
     Returns
     -------
@@ -415,13 +420,15 @@ def plot_mesh_hpc(
     >>> result = plot_mesh_hpc("/hpc/data/grid.nc", use_remote=True)
     >>> open("mesh.png", "wb").write(base64.b64decode(result["png_b64"]))
     """
-    from .plotting import plot_mesh
+    from .plotting import _resolve_plot_paths, plot_mesh
+
+    resolved_grid, _ = _resolve_plot_paths(grid_path, None, session_id, dataset_handle)
 
     def _local() -> Dict[str, Any]:
         import base64
         import json
 
-        items = plot_mesh(grid_path, width=width, height=height)
+        items = plot_mesh(resolved_grid, width=width, height=height)
         # plot_mesh returns [ImageContent, TextContent]; extract png_b64 + metadata
         img = items[0]
         meta = json.loads(items[1].text)
@@ -439,19 +446,19 @@ def _local() -> Dict[str, Any]:
         tool_name="plot_mesh_hpc",
         use_remote=use_remote,
         endpoint=endpoint,
-        path_hint=grid_path,
+        path_hint=resolved_grid,
         session_id=session_id,
         local_call=_local,
         remote_call=lambda agent: _run_sync(
-            lambda: agent.plot_mesh_remote(grid_path, width, height, use_remote)
+            lambda: agent.plot_mesh_remote(resolved_grid, width, height, use_remote)
         ),
     )
     return _plot_result_to_mcp_contents(result)
 
 
 def plot_variable_hpc(
-    grid_path: str,
-    data_path: str,
+    grid_path: str | None = None,
+    data_path: str | None = None,
     variable_name: Optional[str] = None,
     width: int = 800,
     height: int = 400,
@@ -462,15 +469,18 @@ def plot_variable_hpc(
     use_remote: bool = False,
     endpoint: str | None = None,
     session_id: str | None = None,
+    dataset_handle: str | None = None,
 ) -> list[Any]:
     """Render a face-centered variable as a filled-polygon PNG with optional HPC execution.
 
     Parameters
     ----------
-    grid_path : str
-        Path to mesh grid file (local or HPC filesystem).
-    data_path : str
-        Path to data file (local or HPC filesystem).
+    grid_path : str | None
+        Path to mesh grid file (local or HPC filesystem). Optional when
+        ``session_id`` and ``dataset_handle`` are provided.
+    data_path : str | None
+        Path to data file (local or HPC filesystem). Optional when
+        ``session_id`` and ``dataset_handle`` are provided.
     variable_name : str | None
         Variable to plot. If None, the first face-centered variable is used.
     width : int
@@ -487,6 +497,10 @@ def plot_variable_hpc(
         Custom plot title.
     use_remote : bool
         If True and HPC is configured, render on the remote endpoint.
+    session_id, dataset_handle : str | None
+        When both are provided, the grid and data paths are looked up
+        from the registered session dataset instead of ``grid_path`` /
+        ``data_path``.
 
     Returns
     -------
@@ -497,15 +511,19 @@ def plot_variable_hpc(
         - grid_info: {n_face, n_node, n_edge}
         - execution_venue: "local" or "hpc:<endpoint_id>"
     """
-    from .plotting import plot_variable
+    from .plotting import _resolve_plot_paths, plot_variable
+
+    resolved_grid, resolved_data = _resolve_plot_paths(
+        grid_path, data_path, session_id, dataset_handle
+    )
 
     def _local() -> Dict[str, Any]:
         import base64
         import json
 
         items = plot_variable(
-            grid_path,
-            data_path,
+            resolved_grid,
+            resolved_data,
             variable_name,
             width=width,
             height=height,
@@ -531,13 +549,13 @@ def _local() -> Dict[str, Any]:
         tool_name="plot_variable_hpc",
         use_remote=use_remote,
         endpoint=endpoint,
-        path_hint=grid_path,
+        path_hint=resolved_grid,
         session_id=session_id,
         local_call=_local,
         remote_call=lambda agent: _run_sync(
             lambda: agent.plot_variable_remote(
-                grid_path,
-                data_path,
+                resolved_grid,
+                resolved_data,
                 variable_name,
                 width,
                 height,
@@ -553,9 +571,9 @@ def _local() -> Dict[str, Any]:
 
 
 def plot_zonal_mean_hpc(
-    grid_path: str,
-    data_path: str,
-    variable_name: str,
+    grid_path: str | None = None,
+    data_path: str | None = None,
+    variable_name: str | None = None,
     width: int = 800,
     height: int = 400,
     lat_spec: Optional[tuple | float | List] = None,
@@ -565,15 +583,18 @@ def plot_zonal_mean_hpc(
     use_remote: bool = False,
     endpoint: str | None = None,
     session_id: str | None = None,
+    dataset_handle: str | None = None,
 ) -> list[Any]:
     """Render a zonal mean profile PNG with optional HPC execution.
 
     Parameters
     ----------
-    grid_path : str
-        Path to mesh grid file (local or HPC filesystem).
-    data_path : str
-        Path to data file (local or HPC filesystem).
+    grid_path : str | None
+        Path to mesh grid file (local or HPC filesystem). Optional when
+        ``session_id`` and ``dataset_handle`` are provided.
+    data_path : str | None
+        Path to data file (local or HPC filesystem). Optional when
+        ``session_id`` and ``dataset_handle`` are provided.
     variable_name : str
         Face-centered variable to compute and plot.
     width : int
@@ -590,6 +611,10 @@ def plot_zonal_mean_hpc(
         Custom plot title.
     use_remote : bool
         If True and HPC is configured, render on the remote endpoint.
+    session_id, dataset_handle : str | None
+        When both are provided, the grid and data paths are looked up
+        from the registered session dataset instead of ``grid_path`` /
+        ``data_path``.
 
     Returns
     -------
@@ -601,15 +626,19 @@ def plot_zonal_mean_hpc(
         - zonal_mean_values: List of zonal mean values
         - execution_venue: "local" or "hpc:<endpoint_id>"
     """
-    from .plotting import plot_zonal_mean
+    from .plotting import _resolve_plot_paths, plot_zonal_mean
+
+    resolved_grid, resolved_data = _resolve_plot_paths(
+        grid_path, data_path, session_id, dataset_handle
+    )
 
     def _local() -> Dict[str, Any]:
         import base64
         import json
 
         items = plot_zonal_mean(
-            grid_path,
-            data_path,
+            resolved_grid,
+            resolved_data,
             variable_name,
             width=width,
             height=height,
@@ -636,13 +665,13 @@ def _local() -> Dict[str, Any]:
         tool_name="plot_zonal_mean_hpc",
         use_remote=use_remote,
         endpoint=endpoint,
-        path_hint=grid_path,
+        path_hint=resolved_grid,
         session_id=session_id,
         local_call=_local,
         remote_call=lambda agent: _run_sync(
             lambda: agent.plot_zonal_mean_remote(
-                grid_path,
-                data_path,
+                resolved_grid,
+                resolved_data,
                 variable_name,
                 width,
                 height,

tests/test_plotting.py

@@ -426,3 +426,110 @@ def test_plot_zonal_mean_hpc_returns_image_content(self, mock_run):
         metadata = json.loads(result[1].text)
         assert metadata["latitudes"] == [-90.0, 0.0, 90.0]
         assert "png_b64" not in metadata
+
+
+class TestHpcPlotWrappersDatasetHandle:
+    """Issue #25: HPC plot wrappers must accept session_id + dataset_handle and
+    resolve the grid/data paths from the session before dispatching."""
+
+    @staticmethod
+    def _stub_run_result(extra: dict | None = None) -> dict:
+        return {
+            "png_b64": base64.b64encode(b"\x89PNG_fake").decode("utf-8"),
+            "image_size_bytes": 9,
+            "grid_info": {"n_face": 1},
+            "_provenance": {"execution_venue": "local", "warnings": []},
+            **(extra or {}),
+        }
+
+    @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc")
+    def test_plot_mesh_hpc_resolves_dataset_handle(
+        self, mock_run, synthetic_mesh_with_data
+    ):
+        from uxarray_mcp.tools import create_session, register_dataset
+
+        grid_file, _ = synthetic_mesh_with_data
+        session = create_session("plot-mesh-hpc-handle")
+        registered = register_dataset(
+            session_id=session["session_id"],
+            grid_path=grid_file,
+            name="grid-only",
+        )
+
+        mock_run.return_value = self._stub_run_result()
+
+        plot_mesh_hpc(
+            session_id=session["session_id"],
+            dataset_handle=registered["dataset_handle"],
+        )
+
+        kwargs = mock_run.call_args.kwargs
+        assert kwargs["path_hint"] == grid_file
+
+    @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc")
+    def test_plot_variable_hpc_resolves_dataset_handle(
+        self, mock_run, synthetic_mesh_with_data
+    ):
+        from uxarray_mcp.tools import create_session, register_dataset
+
+        grid_file, data_file = synthetic_mesh_with_data
+        session = create_session("plot-variable-hpc-handle")
+        registered = register_dataset(
+            session_id=session["session_id"],
+            grid_path=grid_file,
+            data_path=data_file,
+            name="grid-and-data",
+        )
+
+        mock_run.return_value = self._stub_run_result({"variable_name": "temperature"})
+
+        plot_variable_hpc(
+            variable_name="temperature",
+            session_id=session["session_id"],
+            dataset_handle=registered["dataset_handle"],
+        )
+
+        kwargs = mock_run.call_args.kwargs
+        assert kwargs["path_hint"] == grid_file
+
+    @patch("uxarray_mcp.tools.remote_tools._run_with_optional_hpc")
+    def test_plot_zonal_mean_hpc_resolves_dataset_handle(
+        self, mock_run, synthetic_mesh_with_data
+    ):
+        from uxarray_mcp.tools import create_session, register_dataset
+
+        grid_file, data_file = synthetic_mesh_with_data
+        session = create_session("plot-zonal-hpc-handle")
+        registered = register_dataset(
+            session_id=session["session_id"],
+            grid_path=grid_file,
+            data_path=data_file,
+            name="zonal-handle",
+        )
+
+        mock_run.return_value = self._stub_run_result(
+            {
+                "variable_name": "temperature",
+                "latitudes": [-90.0, 0.0, 90.0],
+                "zonal_mean_values": [270.0, 300.0, 270.0],
+            }
+        )
+
+        plot_zonal_mean_hpc(
+            variable_name="temperature",
+            session_id=session["session_id"],
+            dataset_handle=registered["dataset_handle"],
+        )
+
+        kwargs = mock_run.call_args.kwargs
+        assert kwargs["path_hint"] == grid_file
+
+    def test_plot_mesh_hpc_handle_without_session_raises(self):
+        """dataset_handle without session_id is a clear ValueError."""
+        with pytest.raises(ValueError, match="session_id is required"):
+            plot_mesh_hpc(dataset_handle="some-handle")
+
+    def test_plot_mesh_hpc_no_path_no_handle_raises(self):
+        """At least one of grid_path or dataset_handle must be provided."""
+        with pytest.raises(ValueError, match="grid_path is required"):
+            plot_mesh_hpc()