feat(server): Add hooks.AddOnUnregisterSession functionality (#175)

Add OnUnregisterSession hook functionality to complement the existing 
OnRegisterSession hooks, allowing code to run when a client session 
is being removed from the server. In some cases, the server may want to
do additional work when a session has been closed. For example, in the
SSE server case where you may end up managing various logs for the
duration of the session -- you would want to indicate that the session
was finished.

Author: rwjblue-glean

server/hooks.go

@@ -14,6 +14,8 @@ import (
 // OnRegisterSessionHookFunc is a hook that will be called when a new session is registered.
 type OnRegisterSessionHookFunc func(ctx context.Context, session ClientSession)
 
+// OnUnregisterSessionHookFunc is a hook that will be called when a session is being unregistered.
+type OnUnregisterSessionHookFunc func(ctx context.Context, session ClientSession)
 
 // BeforeAnyHookFunc is a function that is called after the request is
 // parsed but before the method is called.
@@ -63,7 +65,8 @@ type OnAfter{{.HookName}}Func func(ctx context.Context, id any, message *mcp.{{.
 {{end}}
 
 type Hooks struct {
-    OnRegisterSession []OnRegisterSessionHookFunc
+    OnRegisterSession   []OnRegisterSessionHookFunc
+	OnUnregisterSession   []OnUnregisterSessionHookFunc
 	OnBeforeAny      []BeforeAnyHookFunc
 	OnSuccess        []OnSuccessHookFunc
 	OnError          []OnErrorHookFunc
@@ -183,6 +186,19 @@ func (c *Hooks) RegisterSession(ctx context.Context, session ClientSession) {
     }
 }
 
+func (c *Hooks) AddOnUnregisterSession(hook OnUnregisterSessionHookFunc) {
+    c.OnUnregisterSession = append(c.OnUnregisterSession, hook)
+}
+
+func (c *Hooks) UnregisterSession(ctx context.Context, session ClientSession) {
+    if c == nil {
+        return
+    }
+    for _, hook := range c.OnUnregisterSession {
+        hook(ctx, session)
+    }
+}
+
 {{- range .}}
 func (c *Hooks) AddBefore{{.HookName}}(hook OnBefore{{.HookName}}Func) {
 	c.OnBefore{{.HookName}} = append(c.OnBefore{{.HookName}}, hook)

server/internal/gen/hooks.go.tmpl

@@ -206,9 +206,11 @@ func (s *MCPServer) RegisterSession(
 
 // UnregisterSession removes from storage session that is shut down.
 func (s *MCPServer) UnregisterSession(
+	ctx context.Context,
 	sessionID string,
 ) {
-	s.sessions.Delete(sessionID)
+	session, _ := s.sessions.LoadAndDelete(sessionID)
+	s.hooks.UnregisterSession(ctx, session.(ClientSession))
 }
 
 // SendNotificationToAllClients sends a notification to all the currently active clients.

server/server.go

@@ -340,7 +340,6 @@ func TestMCPServer_Tools(t *testing.T) {
 			}`))
 			tt.validate(t, notifications, toolsList.(mcp.JSONRPCMessage))
 		})
-
 	}
 }
 
@@ -794,8 +793,8 @@ func TestMCPServer_HandleInvalidMessages(t *testing.T) {
 			message:     `{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": "invalid"}`,
 			expectedErr: mcp.INVALID_REQUEST,
 			validateErr: func(t *testing.T, err error) {
-				var unparseableErr = &UnparseableMessageError{}
-				var ok = errors.As(err, &unparseableErr)
+				unparseableErr := &UnparseableMessageError{}
+				ok := errors.As(err, &unparseableErr)
 				assert.True(t, ok, "Error should be UnparseableMessageError")
 				assert.Equal(t, mcp.MethodInitialize, unparseableErr.GetMethod())
 				assert.Equal(t, json.RawMessage(`{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": "invalid"}`), unparseableErr.GetMessage())
@@ -1194,7 +1193,6 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
 		assert.Equal(t, "test://something/test-resource/a/b/c", resultContent.URI)
 		assert.Equal(t, "text/plain", resultContent.MIMEType)
 		assert.Equal(t, "test content: something", resultContent.Text)
-
 	})
 }
 
@@ -1422,6 +1420,76 @@ func TestMCPServer_WithHooks(t *testing.T) {
 	assert.IsType(t, afterPingData[0].res, onSuccessData[0].res, "OnSuccess result should be same type as AfterPing result")
 }
 
+func TestMCPServer_SessionHooks(t *testing.T) {
+	var (
+		registerCalled   bool
+		unregisterCalled bool
+
+		registeredContext   context.Context
+		unregisteredContext context.Context
+
+		registeredSession   ClientSession
+		unregisteredSession ClientSession
+	)
+
+	hooks := &Hooks{}
+	hooks.AddOnRegisterSession(func(ctx context.Context, session ClientSession) {
+		registerCalled = true
+		registeredContext = ctx
+		registeredSession = session
+	})
+	hooks.AddOnUnregisterSession(func(ctx context.Context, session ClientSession) {
+		unregisterCalled = true
+		unregisteredContext = ctx
+		unregisteredSession = session
+	})
+
+	server := NewMCPServer(
+		"test-server",
+		"1.0.0",
+		WithHooks(hooks),
+	)
+
+	testSession := &fakeSession{
+		sessionID:           "test-session-id",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 5),
+		initialized:         false,
+	}
+
+	ctx := context.WithoutCancel(context.Background())
+	err := server.RegisterSession(ctx, testSession)
+	require.NoError(t, err)
+
+	assert.True(t, registerCalled, "Register session hook was not called")
+	assert.Equal(t, testSession.SessionID(), registeredSession.SessionID(),
+		"Register hook received wrong session")
+
+	server.UnregisterSession(ctx, testSession.SessionID())
+
+	assert.True(t, unregisterCalled, "Unregister session hook was not called")
+	assert.Equal(t, testSession.SessionID(), unregisteredSession.SessionID(),
+		"Unregister hook received wrong session")
+
+	assert.Equal(t, ctx, unregisteredContext, "Unregister hook received wrong context")
+	assert.Equal(t, ctx, registeredContext, "Register hook received wrong context")
+}
+
+func TestMCPServer_SessionHooks_NilHooks(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0")
+
+	testSession := &fakeSession{
+		sessionID:           "test-session-id",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 5),
+		initialized:         false,
+	}
+
+	ctx := context.WithoutCancel(context.Background())
+	err := server.RegisterSession(ctx, testSession)
+	require.NoError(t, err)
+
+	server.UnregisterSession(ctx, testSession.SessionID())
+}
+
 func TestMCPServer_WithRecover(t *testing.T) {
 	panicToolHandler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 		panic("test panic")

server/server_test.go

@@ -259,7 +259,7 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, fmt.Sprintf("Session registration failed: %v", err), http.StatusInternalServerError)
 		return
 	}
-	defer s.server.UnregisterSession(sessionID)
+	defer s.server.UnregisterSession(r.Context(), sessionID)
 
 	// Start notification handler for this session
 	go func() {
@@ -438,6 +438,7 @@ func (s *SSEServer) SendEventToSession(
 		return fmt.Errorf("event queue full")
 	}
 }
+
 func (s *SSEServer) GetUrlPath(input string) (string, error) {
 	parse, err := url.Parse(input)
 	if err != nil {
@@ -449,6 +450,7 @@ func (s *SSEServer) GetUrlPath(input string) (string, error) {
 func (s *SSEServer) CompleteSseEndpoint() string {
 	return s.baseURL + s.basePath + s.sseEndpoint
 }
+
 func (s *SSEServer) CompleteSsePath() string {
 	path, err := s.GetUrlPath(s.CompleteSseEndpoint())
 	if err != nil {
@@ -460,6 +462,7 @@ func (s *SSEServer) CompleteSsePath() string {
 func (s *SSEServer) CompleteMessageEndpoint() string {
 	return s.baseURL + s.basePath + s.messageEndpoint
 }
+
 func (s *SSEServer) CompleteMessagePath() string {
 	path, err := s.GetUrlPath(s.CompleteMessageEndpoint())
 	if err != nil {

server/sse.go

@@ -204,7 +204,7 @@ func (s *StdioServer) Listen(
 	if err := s.server.RegisterSession(ctx, &stdioSessionInstance); err != nil {
 		return fmt.Errorf("register session: %w", err)
 	}
-	defer s.server.UnregisterSession(stdioSessionInstance.SessionID())
+	defer s.server.UnregisterSession(ctx, stdioSessionInstance.SessionID())
 	ctx = s.server.WithContext(ctx, &stdioSessionInstance)
 
 	// Add in any custom context.