Extract MCP server logic into dedicated package (#1403)

Signed-off-by: Juan Antonio Osorio <[email protected]>

Author: JAORMX

cmd/thv/app/mcp_serve.go

@@ -1,21 +1,28 @@
 package app
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
+	"time"
 
 	"github.com/spf13/cobra"
 
 	s "github.com/stacklok/toolhive/pkg/api"
 	"github.com/stacklok/toolhive/pkg/auth"
+	"github.com/stacklok/toolhive/pkg/logger"
+	mcpserver "github.com/stacklok/toolhive/pkg/mcp/server"
 )
 
 var (
-	host       string
-	port       int
-	enableDocs bool
-	socketPath string
+	host            string
+	port            int
+	enableDocs      bool
+	socketPath      string
+	enableMCPServer bool
+	mcpServerPort   string
+	mcpServerHost   string
 )
 
 var serveCmd = &cobra.Command{
@@ -59,6 +66,39 @@ var serveCmd = &cobra.Command{
 			}
 		}
 
+		// Optionally start MCP server if experimental flag is enabled
+		if enableMCPServer {
+			logger.Info("EXPERIMENTAL: Starting embedded MCP server")
+
+			// Create MCP server configuration
+			mcpConfig := &mcpserver.Config{
+				Host: mcpServerHost,
+				Port: mcpServerPort,
+			}
+
+			// Create and start the MCP server in a goroutine
+			mcpServer, err := mcpserver.New(ctx, mcpConfig)
+			if err != nil {
+				return fmt.Errorf("failed to create MCP server: %w", err)
+			}
+
+			go func() {
+				if err := mcpServer.Start(); err != nil {
+					logger.Errorf("MCP server error: %v", err)
+				}
+			}()
+
+			// Ensure MCP server is shut down on context cancellation
+			go func() {
+				<-ctx.Done()
+				shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer shutdownCancel()
+				if err := mcpServer.Shutdown(shutdownCtx); err != nil {
+					logger.Errorf("Failed to shutdown MCP server: %v", err)
+				}
+			}()
+		}
+
 		return s.Serve(ctx, address, isUnixSocket, debugMode, enableDocs, oidcConfig)
 	},
 }
@@ -71,6 +111,14 @@ func init() {
 	serveCmd.Flags().StringVar(&socketPath, "socket", "", "UNIX socket path to bind the "+
 		"server to (overrides host and port if provided)")
 
+	// Add experimental MCP server flags
+	serveCmd.Flags().BoolVar(&enableMCPServer, "experimental-mcp", false,
+		"EXPERIMENTAL: Enable embedded MCP server for controlling ToolHive")
+	serveCmd.Flags().StringVar(&mcpServerPort, "experimental-mcp-port", mcpserver.DefaultMCPPort,
+		"EXPERIMENTAL: Port for the embedded MCP server")
+	serveCmd.Flags().StringVar(&mcpServerHost, "experimental-mcp-host", "localhost",
+		"EXPERIMENTAL: Host for the embedded MCP server")
+
 	// Add OIDC validation flags
 	AddOIDCFlags(serveCmd)
 }

cmd/thv/app/server.go

@@ -0,0 +1,35 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// getServerLogsArgs holds the arguments for getting server logs
+type getServerLogsArgs struct {
+	Name string `json:"name"`
+}
+
+// GetServerLogs gets logs from a running MCP server
+func (h *Handler) GetServerLogs(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+	// Parse arguments using BindArguments
+	args := &getServerLogsArgs{}
+	if err := request.BindArguments(args); err != nil {
+		return mcp.NewToolResultError(fmt.Sprintf("Failed to parse arguments: %v", err)), nil
+	}
+
+	// Get logs
+	logs, err := h.workloadManager.GetLogs(ctx, args.Name, false)
+	if err != nil {
+		// Check if it's a not found error
+		if strings.Contains(err.Error(), "not found") {
+			return mcp.NewToolResultError(fmt.Sprintf("Server '%s' not found", args.Name)), nil
+		}
+		return mcp.NewToolResultError(fmt.Sprintf("Failed to get server logs: %v", err)), nil
+	}
+
+	return mcp.NewToolResultText(logs), nil
+}

docs/cli/thv_serve.md

@@ -0,0 +1,38 @@
+// Package server provides the MCP (Model Context Protocol) server implementation for ToolHive.
+package server
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/stacklok/toolhive/pkg/registry"
+	"github.com/stacklok/toolhive/pkg/workloads"
+)
+
+// Handler handles MCP tool requests for ToolHive
+type Handler struct {
+	ctx              context.Context
+	workloadManager  workloads.Manager
+	registryProvider registry.Provider
+}
+
+// NewHandler creates a new ToolHive handler
+func NewHandler(ctx context.Context) (*Handler, error) {
+	// Create workload manager
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create workload manager: %w", err)
+	}
+
+	// Create registry provider
+	registryProvider, err := registry.GetDefaultProvider()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get registry provider: %w", err)
+	}
+
+	return &Handler{
+		ctx:              ctx,
+		workloadManager:  workloadManager,
+		registryProvider: registryProvider,
+	}, nil
+}