feat: Add server hooks:OnRequestInitialization (mark3labs#164)

* add hooks:beforeHandleRequest

* Update request_handler.go

use mcp.INVALID_REQUEST replace mcp.PARSE_ERROR

* rename hook name

rename OnBeforeHandleRequest to OnRequestInitialization

* update tmpl file

update tmpl file generate hooks.go/request_handler.go by "go generate"
fix example function name error

* update test

update test message

---------

Co-authored-by: nine <[email protected]>

Author: 2 people

examples/everything/main.go

@@ -44,6 +44,11 @@ func NewMCPServer() *server.MCPServer {
 	hooks.AddBeforeInitialize(func(ctx context.Context, id any, message *mcp.InitializeRequest) {
 		fmt.Printf("beforeInitialize: %v, %v\n", id, message)
 	})
+	hooks.AddOnRequestInitialization(func(ctx context.Context, id any, message any) error {
+		fmt.Printf("AddOnRequestInitialization: %v, %v\n", id, message)
+		// authorization verification and other preprocessing tasks are performed.
+		return nil
+	})
 	hooks.AddAfterInitialize(func(ctx context.Context, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult) {
 		fmt.Printf("afterInitialize: %v, %v, %v\n", id, message, result)
 	})

server/hooks.go

@@ -59,6 +59,11 @@ type OnSuccessHookFunc func(ctx context.Context, id any, method mcp.MCPMethod, m
 // })
 type OnErrorHookFunc func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error)
 
+// OnRequestInitializationFunc is a function that called before handle diff request method
+// Should any errors arise during func execution, the service will promptly return the corresponding error message.
+type OnRequestInitializationFunc func(ctx context.Context, id any, message any) error
+
+
 {{range .}}
 type OnBefore{{.HookName}}Func func(ctx context.Context, id any, message *mcp.{{.ParamType}})
 type OnAfter{{.HookName}}Func func(ctx context.Context, id any, message *mcp.{{.ParamType}}, result *mcp.{{.ResultType}})
@@ -70,6 +75,7 @@ type Hooks struct {
 	OnBeforeAny      []BeforeAnyHookFunc
 	OnSuccess        []OnSuccessHookFunc
 	OnError          []OnErrorHookFunc
+	OnRequestInitialization       []OnRequestInitializationFunc
 {{- range .}}
 	OnBefore{{.HookName}} []OnBefore{{.HookName}}Func
 	OnAfter{{.HookName}}  []OnAfter{{.HookName}}Func
@@ -199,6 +205,23 @@ func (c *Hooks) UnregisterSession(ctx context.Context, session ClientSession) {
     }
 }
 
+func (c *Hooks) AddOnRequestInitialization(hook OnRequestInitializationFunc) {
+	c.OnRequestInitialization = append(c.OnRequestInitialization, hook)
+}
+
+func (c *Hooks) onRequestInitialization(ctx context.Context, id any, message any) error {
+	if c == nil {
+		return nil
+	}
+	for _, hook := range c.OnRequestInitialization {
+		err := hook(ctx, id, message)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 {{- range .}}
 func (c *Hooks) AddBefore{{.HookName}}(hook OnBefore{{.HookName}}Func) {
 	c.OnBefore{{.HookName}} = append(c.OnBefore{{.HookName}}, hook)

server/internal/gen/hooks.go.tmpl

@@ -63,6 +63,15 @@ func (s *MCPServer) HandleMessage(
 		return nil
 	}
 
+	handleErr := s.hooks.onRequestInitialization(ctx, baseMessage.ID, message)
+    if handleErr != nil {
+    	return createErrorResponse(
+    		baseMessage.ID,
+    		mcp.INVALID_REQUEST,
+    		handleErr.Error(),
+    	)
+    }
+
 	switch baseMessage.Method {
 	{{- range .}}
 	case mcp.{{.MethodName}}:

server/internal/gen/request_handler.go.tmpl

@@ -1265,13 +1265,14 @@ var _ ClientSession = fakeSession{}
 func TestMCPServer_WithHooks(t *testing.T) {
 	// Create hook counters to verify calls
 	var (
-		beforeAnyCount   int
-		onSuccessCount   int
-		onErrorCount     int
-		beforePingCount  int
-		afterPingCount   int
-		beforeToolsCount int
-		afterToolsCount  int
+		beforeAnyCount               int
+		onSuccessCount               int
+		onErrorCount                 int
+		beforePingCount              int
+		afterPingCount               int
+		beforeToolsCount             int
+		afterToolsCount              int
+		onRequestInitializationCount int
 	)
 
 	// Collectors for message and result types
@@ -1335,6 +1336,11 @@ func TestMCPServer_WithHooks(t *testing.T) {
 		afterToolsCount++
 	})
 
+	hooks.AddOnRequestInitialization(func(ctx context.Context, id any, message any) error {
+		onRequestInitializationCount++
+		return nil
+	})
+
 	// Create a server with the hooks
 	server := NewMCPServer(
 		"test-server",
@@ -1398,10 +1404,11 @@ func TestMCPServer_WithHooks(t *testing.T) {
 	assert.Equal(t, 1, afterPingCount, "afterPing should be called once")
 	assert.Equal(t, 1, beforeToolsCount, "beforeListTools should be called once")
 	assert.Equal(t, 1, afterToolsCount, "afterListTools should be called once")
-
 	// General hooks should be called for all methods
 	// beforeAny is called for all 4 methods (initialize, ping, tools/list, tools/call)
 	assert.Equal(t, 4, beforeAnyCount, "beforeAny should be called for each method")
+	// onRequestInitialization is called for all 4 methods (initialize, ping, tools/list, tools/call)
+	assert.Equal(t, 4, onRequestInitializationCount, "onRequestInitializationCount should be called for each method")
 	// onSuccess is called for all 3 success methods (initialize, ping, tools/list)
 	assert.Equal(t, 3, onSuccessCount, "onSuccess should be called after all successful invocations")