Skip to content

Commit b7243b8

Browse files
Add session info metadata to initialize response
- Add middleware to enrich InitializeResult with session information - Include user details (get_me response) in authenticated mode - Include enabled toolsets, tools, read-only mode, and lockdown mode - Handle both authenticated and unauthenticated modes appropriately - Add comprehensive unit tests for the new functionality - Remove suggestion to call get_me from auth_login success message Co-authored-by: SamMorrowDrums <4811358+SamMorrowDrums@users.noreply.github.com>
1 parent 590834e commit b7243b8

File tree

3 files changed

+462
-3
lines changed

3 files changed

+462
-3
lines changed

internal/ghmcp/server.go

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ func NewMCPServer(cfg MCPServerConfig) (*mcp.Server, error) {
216216
// Add middlewares
217217
ghServer.AddReceivingMiddleware(addGitHubAPIErrorToContext)
218218
ghServer.AddReceivingMiddleware(addUserAgentsMiddleware(cfg, clients.rest, clients.gqlHTTP))
219+
ghServer.AddReceivingMiddleware(addSessionInfoMiddleware(cfg, clients.rest, enabledToolsets, instructionToolsets))
219220

220221
// Create dependencies for tool handlers
221222
deps := github.NewBaseDeps(
@@ -346,6 +347,14 @@ func NewUnauthenticatedMCPServer(cfg MCPServerConfig) (*UnauthenticatedServerRes
346347
// Add error context middleware
347348
ghServer.AddReceivingMiddleware(addGitHubAPIErrorToContext)
348349

350+
// Add session info middleware for unauthenticated mode
351+
// This will show configuration but no user info until authenticated
352+
instructionToolsets := enabledToolsets
353+
if instructionToolsets == nil {
354+
instructionToolsets = github.GetDefaultToolsetIDs()
355+
}
356+
ghServer.AddReceivingMiddleware(addUnauthenticatedSessionInfoMiddleware(cfg, enabledToolsets, instructionToolsets))
357+
349358
// Create auth tool dependencies with a callback for when auth completes
350359
authDeps := github.AuthToolDependencies{
351360
AuthManager: authManager,
@@ -841,3 +850,156 @@ func addUserAgentsMiddleware(cfg MCPServerConfig, restClient *gogithub.Client, g
841850
}
842851
}
843852
}
853+
854+
// addSessionInfoMiddleware enriches the InitializeResult with session information
855+
// including user details, enabled toolsets, and configuration flags.
856+
func addSessionInfoMiddleware(cfg MCPServerConfig, restClient *gogithub.Client, enabledToolsets []string, instructionToolsets []string) func(next mcp.MethodHandler) mcp.MethodHandler {
857+
return func(next mcp.MethodHandler) mcp.MethodHandler {
858+
return func(ctx context.Context, method string, request mcp.Request) (result mcp.Result, err error) {
859+
// Only intercept initialize method
860+
if method != "initialize" {
861+
return next(ctx, method, request)
862+
}
863+
864+
// Call the next handler to get the InitializeResult
865+
result, err = next(ctx, method, request)
866+
if err != nil {
867+
return result, err
868+
}
869+
870+
// Cast to InitializeResult to add metadata
871+
initResult, ok := result.(*mcp.InitializeResult)
872+
if !ok {
873+
// If we can't cast, just return the original result
874+
return result, err
875+
}
876+
877+
// Build session info metadata
878+
sessionInfo := make(map[string]any)
879+
880+
// Add configuration information
881+
sessionInfo["readOnlyMode"] = cfg.ReadOnly
882+
sessionInfo["lockdownMode"] = cfg.LockdownMode
883+
sessionInfo["dynamicToolsets"] = cfg.DynamicToolsets
884+
885+
// Add toolsets information
886+
switch {
887+
case enabledToolsets == nil:
888+
// nil means "use defaults"
889+
sessionInfo["enabledToolsets"] = instructionToolsets
890+
sessionInfo["toolsetsMode"] = "default"
891+
case len(enabledToolsets) == 0:
892+
sessionInfo["enabledToolsets"] = []string{}
893+
sessionInfo["toolsetsMode"] = "none"
894+
default:
895+
sessionInfo["enabledToolsets"] = enabledToolsets
896+
sessionInfo["toolsetsMode"] = "explicit"
897+
}
898+
899+
// Add enabled tools if specified
900+
if len(cfg.EnabledTools) > 0 {
901+
sessionInfo["enabledTools"] = cfg.EnabledTools
902+
}
903+
904+
// Try to fetch user information (get_me equivalent)
905+
// If it fails, we simply omit it from the session info
906+
if user, _, err := restClient.Users.Get(ctx, ""); err == nil && user != nil {
907+
userInfo := map[string]any{
908+
"login": user.GetLogin(),
909+
"id": user.GetID(),
910+
"profileURL": user.GetHTMLURL(),
911+
"avatarURL": user.GetAvatarURL(),
912+
}
913+
914+
// Add optional fields if they exist
915+
if name := user.GetName(); name != "" {
916+
userInfo["name"] = name
917+
}
918+
if email := user.GetEmail(); email != "" {
919+
userInfo["email"] = email
920+
}
921+
if bio := user.GetBio(); bio != "" {
922+
userInfo["bio"] = bio
923+
}
924+
if company := user.GetCompany(); company != "" {
925+
userInfo["company"] = company
926+
}
927+
if location := user.GetLocation(); location != "" {
928+
userInfo["location"] = location
929+
}
930+
931+
sessionInfo["user"] = userInfo
932+
}
933+
934+
// Set the metadata on the InitializeResult
935+
if initResult.Meta == nil {
936+
initResult.Meta = make(mcp.Meta)
937+
}
938+
initResult.Meta["sessionInfo"] = sessionInfo
939+
940+
return initResult, nil
941+
}
942+
}
943+
}
944+
945+
// addUnauthenticatedSessionInfoMiddleware enriches the InitializeResult with session information
946+
// for unauthenticated servers. This shows configuration but no user info until authenticated.
947+
func addUnauthenticatedSessionInfoMiddleware(cfg MCPServerConfig, enabledToolsets []string, instructionToolsets []string) func(next mcp.MethodHandler) mcp.MethodHandler {
948+
return func(next mcp.MethodHandler) mcp.MethodHandler {
949+
return func(ctx context.Context, method string, request mcp.Request) (result mcp.Result, err error) {
950+
// Only intercept initialize method
951+
if method != "initialize" {
952+
return next(ctx, method, request)
953+
}
954+
955+
// Call the next handler to get the InitializeResult
956+
result, err = next(ctx, method, request)
957+
if err != nil {
958+
return result, err
959+
}
960+
961+
// Cast to InitializeResult to add metadata
962+
initResult, ok := result.(*mcp.InitializeResult)
963+
if !ok {
964+
// If we can't cast, just return the original result
965+
return result, err
966+
}
967+
968+
// Build session info metadata (without user info for unauthenticated mode)
969+
sessionInfo := make(map[string]any)
970+
971+
// Add configuration information
972+
sessionInfo["readOnlyMode"] = cfg.ReadOnly
973+
sessionInfo["lockdownMode"] = cfg.LockdownMode
974+
sessionInfo["dynamicToolsets"] = cfg.DynamicToolsets
975+
sessionInfo["authenticated"] = false
976+
977+
// Add toolsets information
978+
switch {
979+
case enabledToolsets == nil:
980+
// nil means "use defaults"
981+
sessionInfo["enabledToolsets"] = instructionToolsets
982+
sessionInfo["toolsetsMode"] = "default"
983+
case len(enabledToolsets) == 0:
984+
sessionInfo["enabledToolsets"] = []string{}
985+
sessionInfo["toolsetsMode"] = "none"
986+
default:
987+
sessionInfo["enabledToolsets"] = enabledToolsets
988+
sessionInfo["toolsetsMode"] = "explicit"
989+
}
990+
991+
// Add enabled tools if specified
992+
if len(cfg.EnabledTools) > 0 {
993+
sessionInfo["enabledTools"] = cfg.EnabledTools
994+
}
995+
996+
// Set the metadata on the InitializeResult
997+
if initResult.Meta == nil {
998+
initResult.Meta = make(mcp.Meta)
999+
}
1000+
initResult.Meta["sessionInfo"] = sessionInfo
1001+
1002+
return initResult, nil
1003+
}
1004+
}
1005+
}

0 commit comments

Comments
 (0)