diff --git a/jupyter_mcp_server/server.py b/jupyter_mcp_server/server.py index 5c7d6a8..57a1349 100644 --- a/jupyter_mcp_server/server.py +++ b/jupyter_mcp_server/server.py @@ -410,7 +410,7 @@ async def execute_cell( @mcp.tool() async def insert_execute_code_cell( - cell_index: Annotated[int, Field(description="Index of the cell to insert and execute (0-based)", ge=0)], + cell_index: Annotated[int, Field(description="Index of the cell to insert and execute (0-based)", ge=-1)], cell_source: Annotated[str, Field(description="Code source for the cell")], timeout: Annotated[int, Field(description="Maximum seconds to wait for execution")] = 90, ) -> Annotated[list[str | ImageContent], Field(description="List of outputs from the executed cell")]: @@ -743,4 +743,4 @@ async def get_registered_tools(): tools.append(tool_dict) - return tools \ No newline at end of file + return tools diff --git a/jupyter_mcp_server/utils.py b/jupyter_mcp_server/utils.py index be9d560..b6632fc 100644 --- a/jupyter_mcp_server/utils.py +++ b/jupyter_mcp_server/utils.py @@ -5,6 +5,7 @@ import re import asyncio import time +import json from typing import Any, Union from mcp.types import ImageContent from jupyter_mcp_server.config import ALLOW_IMG_OUTPUT @@ -83,7 +84,9 @@ def extract_output(output: Union[dict, Any]) -> Union[str, ImageContent]: return strip_ansi_codes(str(text)) elif output_type in ["display_data", "execute_result"]: - data = output.get("data", {}) + + data = output.get("data", {}) + if "image/png" in data: if ALLOW_IMG_OUTPUT: try: @@ -93,6 +96,8 @@ def extract_output(output: Union[dict, Any]) -> Union[str, ImageContent]: return "[Image Output (PNG) - Error processing image]" else: return "[Image Output (PNG) - Image display disabled]" + + if "text/plain" in data: plain_text = data["text/plain"] if hasattr(plain_text, 'source'): @@ -943,4 +948,4 @@ async def get_notebook_model(serverapp: Any, notebook_path: str): return None nb = NotebookModel() nb._doc = ydoc - return nb \ No newline at end of file + return nb diff --git a/tests/test_tools.py b/tests/test_tools.py index 2ca9acd..3288031 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -108,6 +108,11 @@ async def check_and_delete_cell(client: MCPClient, index, expected_type, content expected_result = eval(code_content) assert int(code_result['result'][0]) == expected_result + # Testing appending code cell to bottom of notebook + code_result = await mcp_client_parametrized.insert_execute_code_cell(-1, code_content) + expected_result = eval(code_content) + assert int(code_result['result'][0]) == expected_result + # Test overwrite_cell_source new_code_content = f"({code_content}) * 2" result = await mcp_client_parametrized.overwrite_cell_source(1, new_code_content) @@ -267,4 +272,4 @@ async def test_list_kernels(mcp_client_parametrized: MCPClient): kernel_list = await mcp_client_parametrized.list_kernels() logging.debug(f"Kernel list: {kernel_list}") # Check for either TSV header or "No kernels found" message - assert "ID\tName\tDisplay_Name\tLanguage\tState\tConnections\tLast_Activity\tEnvironment" in kernel_list \ No newline at end of file + assert "ID\tName\tDisplay_Name\tLanguage\tState\tConnections\tLast_Activity\tEnvironment" in kernel_list