Unified Client Transport Layer for Streamable HTTP Support

client/client.go

@@ -1,84 +1,317 @@
-// Package client provides MCP (Model Control Protocol) client implementations.
 package client
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"sync/atomic"
 
+	"github.com/mark3labs/mcp-go/client/transport"
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
-// MCPClient represents an MCP client interface
-type MCPClient interface {
-	// Initialize sends the initial connection request to the server
-	Initialize(
-		ctx context.Context,
-		request mcp.InitializeRequest,
-	) (*mcp.InitializeResult, error)
-
-	// Ping checks if the server is alive
-	Ping(ctx context.Context) error
-
-	// ListResources requests a list of available resources from the server
-	ListResources(
-		ctx context.Context,
-		request mcp.ListResourcesRequest,
-	) (*mcp.ListResourcesResult, error)
-
-	// ListResourceTemplates requests a list of available resource templates from the server
-	ListResourceTemplates(
-		ctx context.Context,
-		request mcp.ListResourceTemplatesRequest,
-	) (*mcp.ListResourceTemplatesResult,
-		error)
-
-	// ReadResource reads a specific resource from the server
-	ReadResource(
-		ctx context.Context,
-		request mcp.ReadResourceRequest,
-	) (*mcp.ReadResourceResult, error)
-
-	// Subscribe requests notifications for changes to a specific resource
-	Subscribe(ctx context.Context, request mcp.SubscribeRequest) error
-
-	// Unsubscribe cancels notifications for a specific resource
-	Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error
-
-	// ListPrompts requests a list of available prompts from the server
-	ListPrompts(
-		ctx context.Context,
-		request mcp.ListPromptsRequest,
-	) (*mcp.ListPromptsResult, error)
-
-	// GetPrompt retrieves a specific prompt from the server
-	GetPrompt(
-		ctx context.Context,
-		request mcp.GetPromptRequest,
-	) (*mcp.GetPromptResult, error)
-
-	// ListTools requests a list of available tools from the server
-	ListTools(
-		ctx context.Context,
-		request mcp.ListToolsRequest,
-	) (*mcp.ListToolsResult, error)
-
-	// CallTool invokes a specific tool on the server
-	CallTool(
-		ctx context.Context,
-		request mcp.CallToolRequest,
-	) (*mcp.CallToolResult, error)
-
-	// SetLevel sets the logging level for the server
-	SetLevel(ctx context.Context, request mcp.SetLevelRequest) error
-
-	// Complete requests completion options for a given argument
-	Complete(
-		ctx context.Context,
-		request mcp.CompleteRequest,
-	) (*mcp.CompleteResult, error)
-
-	// Close client connection and cleanup resources
-	Close() error
-
-	// OnNotification registers a handler for notifications
-	OnNotification(handler func(notification mcp.JSONRPCNotification))
+// Client implements the MCP client.
+type Client struct {
+	transport transport.Interface
+
+	initialized   bool
+	notifications []func(mcp.JSONRPCNotification)
+	notifyMu      sync.RWMutex
+	requestID     atomic.Int64
+	capabilities  mcp.ServerCapabilities
+}
+
+// NewClient creates a new MCP client with the given transport.
+// Usage:
+//
+//	client, err := NewClient(transport.NewStdio("mcp", nil, "--stdio"))
+//	if err != nil {
+//		log.Fatalf("Failed to create client: %v", err)
+//	}
+func NewClient(transport transport.Interface) *Client {
+	return &Client{
+		transport: transport,
+	}
+}
+
+// Start initiates the connection to the server.
+// Must be called before using the client.
+func (c *Client) Start(ctx context.Context) error {
+	if c.transport == nil {
+		return fmt.Errorf("transport is nil")
+	}
+	err := c.transport.Start(ctx)
+	if err != nil {
+		return err
+	}
+
+	c.transport.SetNotificationHandler(func(notification mcp.JSONRPCNotification) {
+		c.notifyMu.RLock()
+		defer c.notifyMu.RUnlock()
+		for _, handler := range c.notifications {
+			handler(notification)
+		}
+	})
+	return nil
+}
+
+// Close shuts down the client and closes the transport.
+func (c *Client) Close() error {
+	return c.transport.Close()
+}
+
+// OnNotification registers a handler function to be called when notifications are received.
+// Multiple handlers can be registered and will be called in the order they were added.
+func (c *Client) OnNotification(
+	handler func(notification mcp.JSONRPCNotification),
+) {
+	c.notifyMu.Lock()
+	defer c.notifyMu.Unlock()
+	c.notifications = append(c.notifications, handler)
+}
+
+// sendRequest sends a JSON-RPC request to the server and waits for a response.
+// Returns the raw JSON response message or an error if the request fails.
+func (c *Client) sendRequest(
+	ctx context.Context,
+	method string,
+	params interface{},
+) (*json.RawMessage, error) {
+	if !c.initialized && method != "initialize" {
+		return nil, fmt.Errorf("client not initialized")
+	}
+
+	id := c.requestID.Add(1)
+
+	request := transport.JSONRPCRequest{
+		JSONRPC: mcp.JSONRPC_VERSION,
+		ID:      id,
+		Method:  method,
+		Params:  params,
+	}
+
+	response, err := c.transport.SendRequest(ctx, request)
+	if err != nil {
+		return nil, fmt.Errorf("transport error: %w", err)
+	}
+
+	if response.Error != nil {
+		return nil, errors.New(response.Error.Message)
+	}
+
+	return &response.Result, nil
+}
+
+// Initialize negotiates with the server.
+// Must be called after Start, and before any request methods.
+func (c *Client) Initialize(
+	ctx context.Context,
+	request mcp.InitializeRequest,
+) (*mcp.InitializeResult, error) {
+	// Ensure we send a params object with all required fields
+	params := struct {
+		ProtocolVersion string                 `json:"protocolVersion"`
+		ClientInfo      mcp.Implementation     `json:"clientInfo"`
+		Capabilities    mcp.ClientCapabilities `json:"capabilities"`
+	}{
+		ProtocolVersion: request.Params.ProtocolVersion,
+		ClientInfo:      request.Params.ClientInfo,
+		Capabilities:    request.Params.Capabilities, // Will be empty struct if not set
+	}
+
+	response, err := c.sendRequest(ctx, "initialize", params)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.InitializeResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	// Store capabilities
+	c.capabilities = result.Capabilities
+
+	// Send initialized notification
+	notification := mcp.JSONRPCNotification{
+		JSONRPC: mcp.JSONRPC_VERSION,
+		Notification: mcp.Notification{
+			Method: "notifications/initialized",
+		},
+	}
+
+	err = c.transport.SendNotification(ctx, notification)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"failed to send initialized notification: %w",
+			err,
+		)
+	}
+
+	c.initialized = true
+	return &result, nil
+}
+
+func (c *Client) Ping(ctx context.Context) error {
+	_, err := c.sendRequest(ctx, "ping", nil)
+	return err
+}
+
+func (c *Client) ListResources(
+	ctx context.Context,
+	request mcp.ListResourcesRequest,
+) (*mcp.ListResourcesResult, error) {
+	response, err := c.sendRequest(ctx, "resources/list", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.ListResourcesResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	return &result, nil
+}
+
+func (c *Client) ListResourceTemplates(
+	ctx context.Context,
+	request mcp.ListResourceTemplatesRequest,
+) (*mcp.ListResourceTemplatesResult, error) {
+	response, err := c.sendRequest(
+		ctx,
+		"resources/templates/list",
+		request.Params,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.ListResourceTemplatesResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	return &result, nil
+}
+
+func (c *Client) ReadResource(
+	ctx context.Context,
+	request mcp.ReadResourceRequest,
+) (*mcp.ReadResourceResult, error) {
+	response, err := c.sendRequest(ctx, "resources/read", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	return mcp.ParseReadResourceResult(response)
+}
+
+func (c *Client) Subscribe(
+	ctx context.Context,
+	request mcp.SubscribeRequest,
+) error {
+	_, err := c.sendRequest(ctx, "resources/subscribe", request.Params)
+	return err
+}
+
+func (c *Client) Unsubscribe(
+	ctx context.Context,
+	request mcp.UnsubscribeRequest,
+) error {
+	_, err := c.sendRequest(ctx, "resources/unsubscribe", request.Params)
+	return err
+}
+
+func (c *Client) ListPrompts(
+	ctx context.Context,
+	request mcp.ListPromptsRequest,
+) (*mcp.ListPromptsResult, error) {
+	response, err := c.sendRequest(ctx, "prompts/list", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.ListPromptsResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	return &result, nil
+}
+
+func (c *Client) GetPrompt(
+	ctx context.Context,
+	request mcp.GetPromptRequest,
+) (*mcp.GetPromptResult, error) {
+	response, err := c.sendRequest(ctx, "prompts/get", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	return mcp.ParseGetPromptResult(response)
+}
+
+func (c *Client) ListTools(
+	ctx context.Context,
+	request mcp.ListToolsRequest,
+) (*mcp.ListToolsResult, error) {
+	response, err := c.sendRequest(ctx, "tools/list", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.ListToolsResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	return &result, nil
+}
+
+func (c *Client) CallTool(
+	ctx context.Context,
+	request mcp.CallToolRequest,
+) (*mcp.CallToolResult, error) {
+	response, err := c.sendRequest(ctx, "tools/call", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	return mcp.ParseCallToolResult(response)
+}
+
+func (c *Client) SetLevel(
+	ctx context.Context,
+	request mcp.SetLevelRequest,
+) error {
+	_, err := c.sendRequest(ctx, "logging/setLevel", request.Params)
+	return err
+}
+
+func (c *Client) Complete(
+	ctx context.Context,
+	request mcp.CompleteRequest,
+) (*mcp.CompleteResult, error) {
+	response, err := c.sendRequest(ctx, "completion/complete", request.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	var result mcp.CompleteResult
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	return &result, nil
+}
+
+// Helper methods
+
+// GetTransport gives access to the underlying transport layer.
+// Cast it to the specific transport type and obtain the other helper methods.
+func (c *Client) GetTransport() transport.Interface {
+	return c.transport
 }