improve error handling for tools and server methods (#311)

- Fix tool error handling to properly distinguish between protocol
errors and tool execution errors
- Return structured JSON-RPC errors directly for protocol-level issues
- Embed regular tool errors in CallToolResult as per MCP specification
- Improve server error responses for unknown prompts and tools
- Add comprehensive tests for both structured and regular error handling

mcp/tool_test.go

@@ -7,12 +7,16 @@ package mcp
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 )
 
 // testToolHandler is used for type inference in TestNewServerTool.
@@ -132,3 +136,91 @@ func TestUnmarshalSchema(t *testing.T) {
 
 	}
 }
+
+func TestToolErrorHandling(t *testing.T) {
+	// Construct server and add both tools at the top level
+	server := NewServer(testImpl, nil)
+
+	// Create a tool that returns a structured error
+	structuredErrorHandler := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResultFor[any], error) {
+		return nil, &jsonrpc2.WireError{
+			Code:    CodeInvalidParams,
+			Message: "internal server error",
+		}
+	}
+
+	// Create a tool that returns a regular error
+	regularErrorHandler := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResultFor[any], error) {
+		return nil, fmt.Errorf("tool execution failed")
+	}
+
+	AddTool(server, &Tool{Name: "error_tool", Description: "returns structured error"}, structuredErrorHandler)
+	AddTool(server, &Tool{Name: "regular_error_tool", Description: "returns regular error"}, regularErrorHandler)
+
+	// Connect server and client once
+	ct, st := NewInMemoryTransports()
+	_, err := server.Connect(context.Background(), st, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	client := NewClient(testImpl, nil)
+	cs, err := client.Connect(context.Background(), ct, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cs.Close()
+
+	// Test that structured JSON-RPC errors are returned directly
+	t.Run("structured_error", func(t *testing.T) {
+		// Call the tool
+		_, err = cs.CallTool(context.Background(), &CallToolParams{
+			Name:      "error_tool",
+			Arguments: map[string]any{},
+		})
+
+		// Should get the structured error directly
+		if err == nil {
+			t.Fatal("expected error, got nil")
+		}
+
+		var wireErr *jsonrpc2.WireError
+		if !errors.As(err, &wireErr) {
+			t.Fatalf("expected WireError, got %[1]T: %[1]v", err)
+		}
+
+		if wireErr.Code != CodeInvalidParams {
+			t.Errorf("expected error code %d, got %d", CodeInvalidParams, wireErr.Code)
+		}
+	})
+
+	// Test that regular errors are embedded in tool results
+	t.Run("regular_error", func(t *testing.T) {
+		// Call the tool
+		result, err := cs.CallTool(context.Background(), &CallToolParams{
+			Name:      "regular_error_tool",
+			Arguments: map[string]any{},
+		})
+
+		// Should not get an error at the protocol level
+		if err != nil {
+			t.Fatalf("unexpected protocol error: %v", err)
+		}
+
+		// Should get a result with IsError=true
+		if !result.IsError {
+			t.Error("expected IsError=true, got false")
+		}
+
+		// Should have error message in content
+		if len(result.Content) == 0 {
+			t.Error("expected error content, got empty")
+		}
+
+		if textContent, ok := result.Content[0].(*TextContent); !ok {
+			t.Error("expected TextContent")
+		} else if !strings.Contains(textContent.Text, "tool execution failed") {
+			t.Errorf("expected error message in content, got: %s", textContent.Text)
+		}
+	})
+}