Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions jupyter_mcp_server/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")]:
Expand Down Expand Up @@ -743,4 +743,4 @@ async def get_registered_tools():

tools.append(tool_dict)

return tools
return tools
9 changes: 7 additions & 2 deletions jupyter_mcp_server/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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'):
Expand Down Expand Up @@ -943,4 +948,4 @@ async def get_notebook_model(serverapp: Any, notebook_path: str):
return None
nb = NotebookModel()
nb._doc = ydoc
return nb
return nb
7 changes: 6 additions & 1 deletion tests/test_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
assert "ID\tName\tDisplay_Name\tLanguage\tState\tConnections\tLast_Activity\tEnvironment" in kernel_list
Loading