diff --git a/src/claude_agent_sdk/__init__.py b/src/claude_agent_sdk/__init__.py index 4898bc0b..d1ed76e2 100644 --- a/src/claude_agent_sdk/__init__.py +++ b/src/claude_agent_sdk/__init__.py @@ -208,7 +208,13 @@ def create_sdk_mcp_server( - ClaudeAgentOptions: Configuration for using servers with query() """ from mcp.server import Server - from mcp.types import ImageContent, TextContent, Tool + from mcp.types import ( + BlobResourceContents, + EmbeddedResource, + ImageContent, + TextContent, + Tool, + ) # Create MCP server instance server = Server(name, version=version) @@ -278,12 +284,12 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> Any: # Convert result to MCP format # The decorator expects us to return the content, not a CallToolResult # It will wrap our return value in CallToolResult - content: list[TextContent | ImageContent] = [] + content: list[TextContent | ImageContent | EmbeddedResource] = [] if "content" in result: for item in result["content"]: if item.get("type") == "text": content.append(TextContent(type="text", text=item["text"])) - if item.get("type") == "image": + elif item.get("type") == "image": content.append( ImageContent( type="image", @@ -291,6 +297,21 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> Any: mimeType=item["mimeType"], ) ) + elif item.get("type") == "document": + # Convert document to EmbeddedResource with BlobResourceContents + # This preserves document data through MCP for conversion to + # Anthropic document format in query.py + source = item.get("source", {}) + content.append( + EmbeddedResource( + type="resource", + resource=BlobResourceContents( + uri=f"document://{source.get('type', 'base64')}", + mimeType=source.get("media_type", "application/pdf"), + blob=source.get("data", ""), + ), + ) + ) # Return just the content list - the decorator wraps it return content diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index c30fc159..017d68ae 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -486,6 +486,24 @@ async def _handle_sdk_mcp_request( "mimeType": item.mimeType, } ) + elif hasattr(item, "resource") and getattr(item, "type", None) == "resource": + # EmbeddedResource - check if it's a document (PDF, etc.) + resource = item.resource + uri = getattr(resource, "uri", "") + mime_type = getattr(resource, "mimeType", "") + if uri.startswith("document://") or mime_type == "application/pdf": + # Convert EmbeddedResource to Anthropic document format + source_type = uri.replace("document://", "") if uri.startswith("document://") else "base64" + content.append( + { + "type": "document", + "source": { + "type": source_type, + "media_type": mime_type, + "data": getattr(resource, "blob", ""), + }, + } + ) response_data = {"content": content} if hasattr(result.root, "is_error") and result.root.is_error: diff --git a/tests/test_sdk_mcp_integration.py b/tests/test_sdk_mcp_integration.py index d3260073..650e7bf1 100644 --- a/tests/test_sdk_mcp_integration.py +++ b/tests/test_sdk_mcp_integration.py @@ -263,3 +263,72 @@ async def generate_chart(args: dict[str, Any]) -> dict[str, Any]: assert len(tool_executions) == 1 assert tool_executions[0]["name"] == "generate_chart" assert tool_executions[0]["args"]["title"] == "Sales Report" + + +@pytest.mark.asyncio +async def test_document_content_support(): + """Test that tools can return document content (e.g., PDFs) via EmbeddedResource.""" + + # Create sample base64 PDF data (minimal valid PDF) + pdf_data = base64.b64encode( + b"%PDF-1.0\n1 0 obj<>endobj\n" + b"2 0 obj<>endobj\n" + b"3 0 obj<>endobj\n" + b"xref\n0 4\n0000000000 65535 f \n0000000009 00000 n \n" + b"0000000052 00000 n \n0000000101 00000 n \n" + b"trailer<>\nstartxref\n178\n%%EOF" + ).decode("utf-8") + + tool_executions: list[dict[str, Any]] = [] + + @tool("read_document", "Reads a PDF document", {"filename": str}) + async def read_document(args: dict[str, Any]) -> dict[str, Any]: + tool_executions.append({"name": "read_document", "args": args}) + return { + "content": [ + {"type": "text", "text": f"Document: {args['filename']}"}, + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": pdf_data, + }, + }, + ] + } + + server_config = create_sdk_mcp_server( + name="document-test-server", version="1.0.0", tools=[read_document] + ) + + server = server_config["instance"] + call_handler = server.request_handlers[CallToolRequest] + + doc_request = CallToolRequest( + method="tools/call", + params=CallToolRequestParams( + name="read_document", arguments={"filename": "report.pdf"} + ), + ) + result = await call_handler(doc_request) + + # Verify the result contains both text and document (as EmbeddedResource) + assert len(result.root.content) == 2 + + # Check text content + text_content = result.root.content[0] + assert text_content.type == "text" + assert text_content.text == "Document: report.pdf" + + # Check document content (stored as EmbeddedResource with BlobResourceContents) + doc_content = result.root.content[1] + assert doc_content.type == "resource" + assert hasattr(doc_content, "resource") + assert str(doc_content.resource.uri) == "document://base64" + assert doc_content.resource.mimeType == "application/pdf" + assert doc_content.resource.blob == pdf_data + + # Verify tool execution + assert len(tool_executions) == 1 + assert tool_executions[0]["name"] == "read_document"