@@ -18,6 +18,7 @@ import (
1818 "github.com/hashicorp/terraform-mcp-server/pkg/toolsets"
1919 "github.com/hashicorp/terraform-mcp-server/version"
2020
21+ "github.com/mark3labs/mcp-go/mcp"
2122 "github.com/mark3labs/mcp-go/server"
2223 log "github.com/sirupsen/logrus"
2324 "github.com/spf13/cobra"
@@ -69,6 +70,23 @@ func NewServer(version string, logger *log.Logger, enabledToolsets []string, opt
6970 hooks .AddOnUnregisterSession (func (ctx context.Context , session server.ClientSession ) {
7071 client .EndSessionHandler (ctx , session , logger )
7172 })
73+ // When running multiple sessions of the MCP server (load balancing), calling client.NewSessionHandler
74+ // in both BeforeListTools and BeforeCallTool ensures that a session that was not initialized during
75+ // registration (e.g., due to being routed to a different instance) will still have its clients created
76+ // before any tool calls are made. This provides a safety net to ensure that all sessions have
77+ // the necessary clients initialized regardless of how they are routed.
78+ hooks .AddBeforeListTools (func (ctx context.Context , id any , message * mcp.ListToolsRequest ) {
79+ session := server .ClientSessionFromContext (ctx )
80+ if session != nil {
81+ client .NewSessionHandler (ctx , session , logger )
82+ }
83+ })
84+ hooks .AddBeforeCallTool (func (ctx context.Context , id any , message * mcp.CallToolRequest ) {
85+ session := server .ClientSessionFromContext (ctx )
86+ if session != nil {
87+ client .NewSessionHandler (ctx , session , logger )
88+ }
89+ })
7290
7391 // Add hooks to options
7492 opts = append (opts , server .WithHooks (hooks ))
@@ -210,19 +228,6 @@ func shouldUseStreamableHTTPMode() bool {
210228 os .Getenv ("MCP_ENDPOINT" ) != ""
211229}
212230
213- // shouldUseStatelessMode returns true if the MCP_SESSION_MODE environment variable is set to "stateless"
214- func shouldUseStatelessMode () bool {
215- mode := strings .ToLower (os .Getenv ("MCP_SESSION_MODE" ))
216-
217- // Explicitly check for "stateless" value
218- if mode == "stateless" {
219- return true
220- }
221-
222- // All other values (including empty string, "stateful", or any other value) default to stateful mode
223- return false
224- }
225-
226231// getHTTPPort returns the port from environment variables or default
227232func getHTTPPort () string {
228233 if port := os .Getenv ("TRANSPORT_PORT" ); port != "" {
@@ -239,6 +244,19 @@ func getHTTPHost() string {
239244 return "127.0.0.1"
240245}
241246
247+ // shouldUseStatelessMode returns true if the MCP_SESSION_MODE environment variable is set to "stateless"
248+ func shouldUseStatelessMode () bool {
249+ mode := strings .ToLower (os .Getenv ("MCP_SESSION_MODE" ))
250+
251+ // Explicitly check for "stateless" value
252+ if mode == "stateless" {
253+ return true
254+ }
255+
256+ // All other values (including empty string, "stateful", or any other value) default to stateful mode
257+ return false
258+ }
259+
242260// Add function to get endpoint path from environment or flag
243261func getEndpointPath (cmd * cobra.Command ) string {
244262 // First check environment variable
0 commit comments