diff --git a/README.md b/README.md index 7dbee3f..a898227 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,20 @@ You can also filter resources to mask sensitive information: {{ support_ticket_info | regex_replace(r'\b[\w.+-]+@[\w.-]+\.\w+\b', '[EMAIL_REDACTED]') }} ``` +### Interactively selecting MCP resources and tools + +Template variables can be given values either directly or by selecting an MCP tool to call. For example: + +``` +# Spec Template (spec.md) + +## Task description +{{ task }} +``` + +Running create-spec will then prompt you to either provide a direct value for `task` or select an MCP server, tool and args to call to fetch the task description. + + ### Initialize a project ``` cxk init diff --git a/engine/globals.py b/engine/globals.py index cf4ce30..b6934a4 100644 --- a/engine/globals.py +++ b/engine/globals.py @@ -5,12 +5,13 @@ from mcp.shared.metadata_utils import get_display_name from pydantic import AnyUrl -from mcp_client import get_client_session_by_server +from mcp_client import get_client_session_by_server, handle_binary_content from prompt import PromptHelper +from state import State from util.parse import parse_input_string -def create_mcp_tool_function(prompt_helper: PromptHelper): +def create_mcp_tool_function(prompt_helper: PromptHelper, state: State): """Create a call_mcp_tool function with state bound.""" async def call_mcp_tool(server: str, tool_name: str, args: dict) -> str | dict[str, Any]: @@ -26,6 +27,13 @@ async def call_mcp_tool(server: str, tool_name: str, args: dict) -> str | dict[s result_unstructured = result.content[0] if isinstance(result_unstructured, types.TextContent): return parse_input_string(result_unstructured.text) + elif isinstance(result_unstructured, types.ImageContent | types.AudioContent): + if state.config_dir: + file_path = handle_binary_content(state.config_dir, result_unstructured) + return file_path if file_path else "" + else: + logging.error("Config directory not available for binary data storage") + return "" else: return "" except Exception as e: @@ -38,7 +46,7 @@ async def call_mcp_tool(server: str, tool_name: str, args: dict) -> str | dict[s return call_mcp_tool -def create_mcp_resource_function(): +def create_mcp_resource_function(state: State): """Create a get_mcp_resource function with state bound.""" async def get_mcp_resource(server: str, resource_uri: str) -> str | dict[str, Any]: @@ -50,6 +58,13 @@ async def get_mcp_resource(server: str, resource_uri: str) -> str | dict[str, An result_unstructured = resource.contents[0] if isinstance(result_unstructured, types.TextResourceContents): return parse_input_string(result_unstructured.text) + elif isinstance(result_unstructured, types.BlobResourceContents): + if state.config_dir: + file_path = handle_binary_content(state.config_dir, result_unstructured) + return file_path if file_path else "" + else: + logging.error("Config directory not available for binary data storage") + return "" else: logging.debug(f"Resource {resource_uri} returned non-text content") return "" diff --git a/engine/template_engine.py b/engine/template_engine.py index f5d2e4e..966f50b 100644 --- a/engine/template_engine.py +++ b/engine/template_engine.py @@ -31,8 +31,8 @@ def __init__( self._prompt_helper = prompt_helper # Add global functions to env to support MCP tools and resources - self.env.globals["call_tool"] = create_mcp_tool_function(self._prompt_helper) - self.env.globals["get_resource"] = create_mcp_resource_function() + self.env.globals["call_tool"] = create_mcp_tool_function(self._prompt_helper, self._state) + self.env.globals["get_resource"] = create_mcp_resource_function(self._state) @classmethod def from_file(cls, path: str | Path, state: State, prompt_helper: PromptHelper) -> "TemplateEngine": diff --git a/mcp_client/__init__.py b/mcp_client/__init__.py index d5b51e5..b9e04a2 100644 --- a/mcp_client/__init__.py +++ b/mcp_client/__init__.py @@ -1,5 +1,6 @@ """MCP (Model Context Protocol) client implementation.""" +from .binary_data_handler import handle_binary_content, save_binary_data_to_file from .client_session_provider import get_client_session_by_server from .config import HTTPServerConfig, MCPServersConfig, SSEServerConfig, StdioServerConfig from .session_manager import MCPSessionManager, get_session_manager @@ -14,4 +15,6 @@ "KeychainTokenStorageWithFallback", "get_session_manager", "get_client_session_by_server", + "handle_binary_content", + "save_binary_data_to_file", ] diff --git a/mcp_client/binary_data_handler.py b/mcp_client/binary_data_handler.py new file mode 100644 index 0000000..1e18cef --- /dev/null +++ b/mcp_client/binary_data_handler.py @@ -0,0 +1,76 @@ +import base64 +import logging +import mimetypes +from pathlib import Path +from uuid import uuid4 + +from mcp import types + + +def save_binary_data_to_file(config_dir: Path, data: str, mime_type: str, file_prefix: str = "binary") -> str: + """ + Save binary data (base64-encoded) to a file in the config directory. + + Args: + config_dir: The .cxk config directory path + data: Base64-encoded binary data + mime_type: MIME type of the data (e.g., 'image/png', 'audio/wav') + file_prefix: Prefix for the generated filename + + Returns: + Relative path to the saved file from the project root + + Raises: + Exception: If file save fails (errors are logged but not re-raised) + """ + try: + # Create files directory if it doesn't exist + files_dir = config_dir / "files" + files_dir.mkdir(exist_ok=True) + + # Generate unique filename with appropriate extension + extension = mimetypes.guess_extension(mime_type) or "" + filename = f"{file_prefix}_{uuid4().hex[:8]}{extension}" + file_path = files_dir / filename + + # Decode base64 data and write to file + binary_data = base64.b64decode(data) + with open(file_path, "wb") as f: + f.write(binary_data) + + # Return relative path from project root + relative_path = file_path.relative_to(config_dir.parent) + logging.info(f"Saved binary data to {relative_path}") + return str(relative_path) + + except Exception as e: + error_msg = f"Failed to save binary data: {e}" + logging.error(error_msg) + raise Exception(error_msg) from e + + +def handle_binary_content(config_dir: Path, content) -> str: + """ + Handle binary content from MCP responses, saving to file and returning path. + + Args: + config_dir: The .cxk config directory path + content: MCP content object (ImageContent, AudioContent, or BlobResourceContents) + + Returns: + Relative path to the saved file or empty string if handling fails + """ + try: + if isinstance(content, types.ImageContent): + return save_binary_data_to_file(config_dir, content.data, content.mimeType, "image") + elif isinstance(content, types.AudioContent): + return save_binary_data_to_file(config_dir, content.data, content.mimeType, "audio") + elif isinstance(content, types.BlobResourceContents): + # For blob resources, we don't have mime type info, so we'll save as generic file + return save_binary_data_to_file(config_dir, content.blob, "application/octet-stream", "blob") + else: + logging.warning(f"Unsupported binary content type: {type(content)}") + return "" + except Exception as e: + logging.error(f"Error handling binary content: {e}") + return "" diff --git a/prompt/prompt_helper.py b/prompt/prompt_helper.py index 4a210e0..82bd3da 100644 --- a/prompt/prompt_helper.py +++ b/prompt/prompt_helper.py @@ -1,8 +1,9 @@ import logging import questionary +from mcp import types -from mcp_client import get_session_manager +from mcp_client import get_session_manager, handle_binary_content from state import State @@ -221,11 +222,20 @@ async def _collect_var_value_from_mcp(self, var_name: str) -> str: # Handle list of content items content_parts = [] for item in result.content: - # Try to get text content first (safely) - text_content = getattr(item, "text", None) - if text_content: - content_parts.append(str(text_content)) - # For non-text content, convert to string + # Handle binary content types + if isinstance(item, types.ImageContent | types.AudioContent): + if self._state.config_dir: + file_path = handle_binary_content(self._state.config_dir, item) + if file_path: + content_parts.append(file_path) + else: + logging.error(f"Failed to save binary content for variable '{var_name}'") + else: + logging.error("Config directory not available for binary data storage") + # Handle text content + elif isinstance(item, types.TextContent): + content_parts.append(str(item.text)) + # For other content types, convert to string else: content_parts.append(str(item)) content = "\n".join(content_parts) @@ -279,6 +289,6 @@ async def _select_mcp_tool(self, tools) -> str | None: choice_title = f"{tool.name} - {description}" choices.append(questionary.Choice(choice_title, tool.name)) - tool_name = await questionary.select("Select a tool:", choices=choices).ask_async() + tool_name = await questionary.select("Select a tool:", choices=choices, use_search_filter=True).ask_async() return tool_name diff --git a/tests/files/image1.avif b/tests/files/image1.avif new file mode 100644 index 0000000..d330869 Binary files /dev/null and b/tests/files/image1.avif differ diff --git a/tests/mcp_test_server.py b/tests/mcp_test_server.py index 987be3e..3527b0f 100644 --- a/tests/mcp_test_server.py +++ b/tests/mcp_test_server.py @@ -4,7 +4,7 @@ from typing import Any -from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp import FastMCP, Image # Create an MCP server mcp = FastMCP("Test-Server", "1.0.0") @@ -32,3 +32,11 @@ def jsonTest(cloudId: str, ticketId: str, optional_other: str | None = None) -> def get_greeting(name: str) -> str: """Get a personalized greeting""" return f"Hello, {name}!" + + +# Add a binary tool (tests/files/image1.avif) +@mcp.tool() +def get_blob() -> Image: + """Get a binary blob test tool""" + with open("tests/files/image1.avif", "rb") as f: + return Image(data=f.read(), format="avif") diff --git a/tests/templates/spec10.md b/tests/templates/spec10.md new file mode 100644 index 0000000..6b41474 --- /dev/null +++ b/tests/templates/spec10.md @@ -0,0 +1,7 @@ +# Task Template + +## Image links + +{% set image = call_tool('test-mcp', 'get_blob', {}) %} + +{{ image }} \ No newline at end of file diff --git a/tests/templates/spec9.md b/tests/templates/spec9.md index 7adae2f..d361e03 100644 --- a/tests/templates/spec9.md +++ b/tests/templates/spec9.md @@ -1,6 +1,6 @@ # Task Template -## Issue description +## Image links {% set issue = call_tool('github', 'get_issue', {'issue_number': 17, 'owner': 'eyalzh', 'repo': 'browser-control-mcp'}) %} diff --git a/tests/test_binary_data_handler.py b/tests/test_binary_data_handler.py new file mode 100644 index 0000000..95776da --- /dev/null +++ b/tests/test_binary_data_handler.py @@ -0,0 +1,160 @@ +"""Tests for binary data handling functionality.""" + +import base64 +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import Mock + +import pytest +from mcp import types +from pydantic import AnyUrl + +from mcp_client.binary_data_handler import handle_binary_content, save_binary_data_to_file + + +def test_save_binary_data_to_file(): + """Test saving binary data to file.""" + # Create a temporary directory + with TemporaryDirectory() as temp_dir: + config_dir = Path(temp_dir) / ".cxk" + config_dir.mkdir() + + # Create some test binary data + test_data = b"test binary data" + encoded_data = base64.b64encode(test_data).decode("utf-8") + + # Save the data + file_path = save_binary_data_to_file(config_dir, encoded_data, "image/png", "test") + + # Verify the file was created and contains correct data + full_path = Path(temp_dir) / file_path + assert full_path.exists() + + with open(full_path, "rb") as f: + saved_data = f.read() + + assert saved_data == test_data + assert file_path.startswith(".cxk/files/test_") + assert file_path.endswith(".png") + + +def test_handle_image_content(): + """Test handling ImageContent.""" + with TemporaryDirectory() as temp_dir: + config_dir = Path(temp_dir) / ".cxk" + config_dir.mkdir() + + # Create mock ImageContent + test_data = b"fake image data" + encoded_data = base64.b64encode(test_data).decode("utf-8") + + image_content = types.ImageContent( + type="image", + data=encoded_data, + mimeType="image/jpeg" + ) + + # Handle the content + file_path = handle_binary_content(config_dir, image_content) + + # Verify result + assert file_path + assert file_path.startswith(".cxk/files/image_") + assert file_path.endswith(".jpg") # mimetypes maps image/jpeg to .jpg + + # Verify file contents + full_path = Path(temp_dir) / file_path + with open(full_path, "rb") as f: + saved_data = f.read() + assert saved_data == test_data + + +def test_handle_audio_content(): + """Test handling AudioContent.""" + with TemporaryDirectory() as temp_dir: + config_dir = Path(temp_dir) / ".cxk" + config_dir.mkdir() + + # Create mock AudioContent + test_data = b"fake audio data" + encoded_data = base64.b64encode(test_data).decode("utf-8") + + audio_content = types.AudioContent( + type="audio", + data=encoded_data, + mimeType="audio/mpeg" # This should map to .mp3 + ) + + # Handle the content + file_path = handle_binary_content(config_dir, audio_content) + + # Verify result + assert file_path + assert file_path.startswith(".cxk/files/audio_") + # mimetypes may not always find an extension, so just check prefix + assert "audio_" in file_path + + # Verify file contents + full_path = Path(temp_dir) / file_path + with open(full_path, "rb") as f: + saved_data = f.read() + assert saved_data == test_data + + +def test_handle_blob_resource_contents(): + """Test handling BlobResourceContents.""" + with TemporaryDirectory() as temp_dir: + config_dir = Path(temp_dir) / ".cxk" + config_dir.mkdir() + + # Create mock BlobResourceContents + test_data = b"fake blob data" + encoded_data = base64.b64encode(test_data).decode("utf-8") + + blob_content = types.BlobResourceContents( + blob=encoded_data, + uri=AnyUrl("test://blob"), + mimeType="application/octet-stream" + ) + + # Handle the content + file_path = handle_binary_content(config_dir, blob_content) + + # Verify result + assert file_path + assert file_path.startswith(".cxk/files/blob_") + + # Verify file contents + full_path = Path(temp_dir) / file_path + with open(full_path, "rb") as f: + saved_data = f.read() + assert saved_data == test_data + + +def test_handle_unsupported_content(): + """Test handling unsupported content types.""" + with TemporaryDirectory() as temp_dir: + config_dir = Path(temp_dir) / ".cxk" + config_dir.mkdir() + + # Create mock unsupported content + unsupported_content = Mock() + + # Handle the content + file_path = handle_binary_content(config_dir, unsupported_content) + + # Should return empty string for unsupported types + assert file_path == "" + + + +def test_save_binary_data_error_handling(): + """Test error handling when save fails.""" + # Use a non-existent directory to trigger an error + config_dir = Path("/nonexistent/path") + + # This should raise an exception + with pytest.raises(Exception) as exc_info: + save_binary_data_to_file(config_dir, "dGVzdA==", "image/png") + + assert "Failed to save binary data" in str(exc_info.value)