Session store is not a singleton

Will allow us to swap in different session stores

GitOrigin-RevId: 1c2470b

Author: wendorf

cmd/server.go

@@ -36,7 +36,8 @@ func Serve(transport string) *server.MCPServer {
 	s.AddTools(logs.Tools(c)...)
 
 	if transport == "http" {
-		if err := server.NewStreamableHTTPServer(s, server.WithHTTPContextFunc(session.ContextWithHTTPSession)).Start(":10000"); err != nil {
+		sessionStore := session.NewInMemoryStore()
+		if err := server.NewStreamableHTTPServer(s, server.WithHTTPContextFunc(session.ContextWithHTTPSession(sessionStore))).Start(":10000"); err != nil {
 			log.Fatalf("Starting Streamable server: %v\n:", err)
 		}
 	} else {

pkg/session/http.go

@@ -12,12 +12,10 @@ type sessionCtxKeyType struct{}
 
 var sessionCtxKey sessionCtxKeyType
 
-func ContextWithHTTPSession(ctx context.Context, _ *http.Request) context.Context {
-	cs := server.ClientSessionFromContext(ctx)
-	if _, ok := inMemoryDataSingleton[cs.SessionID()]; !ok {
-		inMemoryDataSingleton[cs.SessionID()] = &HTTPSession{}
+func ContextWithHTTPSession(store Store) func(ctx context.Context, _ *http.Request) context.Context {
+	return func(ctx context.Context, _ *http.Request) context.Context {
+		return context.WithValue(ctx, sessionCtxKey, store.Get(server.ClientSessionFromContext(ctx).SessionID()))
 	}
-	return context.WithValue(ctx, sessionCtxKey, inMemoryDataSingleton[cs.SessionID()])
 }
 
 type HTTPSession struct {
@@ -38,4 +36,25 @@ func (h *HTTPSession) SetWorkspace(s string) error {
 	return nil
 }
 
-var inMemoryDataSingleton = map[string]*HTTPSession{}
+type Store interface {
+	Get(sessionID string) *HTTPSession
+}
+
+type inMemoryStore struct {
+	sessions map[string]*HTTPSession
+}
+
+var _ Store = (*inMemoryStore)(nil)
+
+func NewInMemoryStore() Store {
+	return &inMemoryStore{
+		sessions: make(map[string]*HTTPSession),
+	}
+}
+
+func (i *inMemoryStore) Get(sessionID string) *HTTPSession {
+	if _, ok := i.sessions[sessionID]; !ok {
+		i.sessions[sessionID] = &HTTPSession{}
+	}
+	return i.sessions[sessionID]
+}

pkg/session/http_test.go

@@ -10,9 +10,11 @@ import (
 )
 
 func TestHTTPSession(t *testing.T) {
+	sessionStore := session.NewInMemoryStore()
+	contextWithHTTPSession := session.ContextWithHTTPSession(sessionStore)
 	{
 		ctxOne := (&server.MCPServer{}).WithContext(context.Background(), fakeSession{sessionID: "one"})
-		ctxOne = session.ContextWithHTTPSession(ctxOne, nil)
+		ctxOne = contextWithHTTPSession(ctxOne, nil)
 
 		sessionOne := session.FromContext(ctxOne)
 
@@ -36,7 +38,7 @@ func TestHTTPSession(t *testing.T) {
 
 	{
 		ctxOneAgain := (&server.MCPServer{}).WithContext(context.Background(), fakeSession{sessionID: "one"})
-		ctxOneAgain = session.ContextWithHTTPSession(ctxOneAgain, nil)
+		ctxOneAgain = contextWithHTTPSession(ctxOneAgain, nil)
 		sessionOneAgain := session.FromContext(ctxOneAgain)
 
 		workspace, err := sessionOneAgain.GetWorkspace()
@@ -50,7 +52,7 @@ func TestHTTPSession(t *testing.T) {
 
 	{
 		ctxTwo := (&server.MCPServer{}).WithContext(context.Background(), fakeSession{sessionID: "two"})
-		ctxTwo = session.ContextWithHTTPSession(ctxTwo, nil)
+		ctxTwo = contextWithHTTPSession(ctxTwo, nil)
 
 		sessionTwo := session.FromContext(ctxTwo)
 
@@ -67,7 +69,7 @@ func TestHTTPSession(t *testing.T) {
 	{
 
 		ctxOneFinal := (&server.MCPServer{}).WithContext(context.Background(), fakeSession{sessionID: "one"})
-		ctxOneFinal = session.ContextWithHTTPSession(ctxOneFinal, nil)
+		ctxOneFinal = contextWithHTTPSession(ctxOneFinal, nil)
 		sessionOneFinal := session.FromContext(ctxOneFinal)
 
 		workspace, err := sessionOneFinal.GetWorkspace()