feat(sse)!: refactor cmd hierarchy to group transports under serve

ccamel · ccamel · commit 1c7d13bbfafa · 2025-04-04T20:16:02.000+02:00
diff --git a/cmd/root.go b/cmd/root.go
@@ -22,3 +22,7 @@ func Execute() {
 		os.Exit(1)
 	}
 }
+
+func init() {
+	cobra.EnableTraverseRunHooks = true
+}
diff --git a/cmd/serve-sse.go b/cmd/serve-sse.go
@@ -0,0 +1,103 @@
+package cmd
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"axone-protocol/axone-mcp/internal/mcp"
+
+	"github.com/justinas/alice"
+
+	"github.com/mark3labs/mcp-go/server"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/hlog"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+)
+
+const (
+	FlagBaseURL    = "base-url"
+	FlagListenAddr = "listen-addr"
+)
+
+const (
+	ReadHeaderTimeout = 5 * time.Second
+)
+
+var (
+	baseURL    string
+	listenAddr string
+)
+
+var serveSseCmd = &cobra.Command{
+	Use:   "sse",
+	Short: "Serve the MCP over SSE (server-sent events)",
+	Long: `Start the MCP server using Server-Sent Events (SSE) to enable streaming over HTTP.
+Typically used for browser-based or reactive clients.`,
+	RunE: func(_ *cobra.Command, _ []string) error {
+		s, err := mcp.NewServer()
+		if err != nil {
+			return err
+		}
+
+		sseServer := server.NewSSEServer(s)
+		chain := loggerChain().Then(sseServer)
+		httpSrv := &http.Server{
+			Addr:              listenAddr,
+			ReadHeaderTimeout: ReadHeaderTimeout,
+			Handler:           chain,
+		}
+
+		go func() {
+			log.Logger.Info().
+				Str("transport", "sse").
+				Str("base_url", baseURL).
+				Str("addr", listenAddr).
+				Str("message_path", sseServer.CompleteMessagePath()).
+				Str("sse_path", sseServer.CompleteSsePath()).
+				Msg("ready")
+			if err := httpSrv.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
+				log.Fatal().Err(err).Msg("failed to start server")
+			}
+		}()
+
+		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+		defer stop()
+
+		<-ctx.Done()
+		log.Info().Msg("shutdown signal received")
+
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		return sseServer.Shutdown(shutdownCtx)
+	},
+}
+
+func loggerChain() alice.Chain {
+	return alice.New(hlog.NewHandler(log.Logger),
+		hlog.AccessHandler(func(r *http.Request, status, size int, duration time.Duration) {
+			zerolog.Ctx(r.Context()).
+				Debug().
+				Str("client", r.RemoteAddr).
+				Str("method", r.Method).
+				Stringer("path", r.URL).
+				Str("user_agent", r.Header.Get("User-Agent")).
+				Int("status", status).
+				Int("size", size).
+				Dur("duration", duration).
+				Msg("")
+		}),
+	)
+}
+
+func init() {
+	serveSseCmd.PersistentFlags().StringVar(&baseURL, FlagBaseURL, "", "The server's base URL")
+	serveSseCmd.PersistentFlags().StringVar(&listenAddr, FlagListenAddr, "127.0.0.1:8081", "The server's listen address")
+
+	serveCmd.AddCommand(serveSseCmd)
+}
diff --git a/cmd/serve.go b/cmd/serve.go
@@ -1,103 +1,24 @@
 package cmd
 
 import (
-	"context"
-	"errors"
-	"net/http"
-	"os"
-	"os/signal"
-	"syscall"
-	"time"
-
-	"axone-protocol/axone-mcp/internal/mcp"
-
-	"github.com/justinas/alice"
-
-	"github.com/mark3labs/mcp-go/server"
-	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/hlog"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 )
 
-const (
-	FlagBaseURL    = "base-url"
-	FlagListenAddr = "listen-addr"
-)
-
-const (
-	ReadHeaderTimeout = 5 * time.Second
-)
-
-var (
-	baseURL    string
-	listenAddr string
-)
-
-var serveSseCmd = &cobra.Command{
-	Use:   "serve-sse",
-	Short: "Start the MCP server (SSE)",
-	RunE: func(_ *cobra.Command, _ []string) error {
-		log.Logger.Info().Msg("starting server")
-
-		s, err := mcp.NewServer()
-		if err != nil {
-			return err
-		}
-
-		sseServer := server.NewSSEServer(s)
-		chain := loggerChain().Then(sseServer)
-		httpSrv := &http.Server{
-			Addr:              listenAddr,
-			ReadHeaderTimeout: ReadHeaderTimeout,
-			Handler:           chain,
-		}
-
-		go func() {
-			log.Logger.Info().
-				Str("transport", "sse").
-				Str("base_url", baseURL).
-				Str("addr", listenAddr).
-				Str("message_path", sseServer.CompleteMessagePath()).
-				Str("sse_path", sseServer.CompleteSsePath()).
-				Msg("listening for connections")
-			if err := httpSrv.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
-				log.Fatal().Err(err).Msg("failed to start server")
-			}
-		}()
-
-		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-		defer stop()
-
-		<-ctx.Done()
-		log.Info().Msg("shutdown signal received")
-
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
-		return sseServer.Shutdown(shutdownCtx)
+// serveCmd represents the base serve command.
+var serveCmd = &cobra.Command{
+	Use:   "serve",
+	Short: "Serve the MCP using a specific transport",
+	Long: `Start the Axone MCP server using the chosen transport:
+SSE for web clients, stdio for command-line and local integrations.`,
+	PersistentPreRun: func(_ *cobra.Command, _ []string) {
+		log.Logger.Info().Msg("starting server...")
+	},
+	PersistentPostRun: func(_ *cobra.Command, _ []string) {
+		log.Logger.Info().Msg("server stopped")
 	},
-}
-
-func loggerChain() alice.Chain {
-	return alice.New(hlog.NewHandler(log.Logger),
-		hlog.AccessHandler(func(r *http.Request, status, size int, duration time.Duration) {
-			zerolog.Ctx(r.Context()).
-				Debug().
-				Str("client", r.RemoteAddr).
-				Str("method", r.Method).
-				Stringer("path", r.URL).
-				Str("user_agent", r.Header.Get("User-Agent")).
-				Int("status", status).
-				Int("size", size).
-				Dur("duration", duration).
-				Msg("")
-		}),
-	)
 }
 
 func init() {
-	serveSseCmd.PersistentFlags().StringVar(&baseURL, FlagBaseURL, "", "The server's base URL")
-	serveSseCmd.PersistentFlags().StringVar(&listenAddr, FlagListenAddr, "127.0.0.1:8081", "The server's listen address")
-
-	rootCmd.AddCommand(serveSseCmd)
+	rootCmd.AddCommand(serveCmd)
 }

Original file line number	Diff line number	Diff line change
`@@ -22,3 +22,7 @@ func Execute() {`
`22`	`22`	`os.Exit(1)`
`23`	`23`	`}`
`24`	`24`	`}`
	`25`	`+`
	`26`	`+func init() {`
	`27`	`+ cobra.EnableTraverseRunHooks = true`
	`28`	`+}`