From fcb86b1056a32fb10cee5e26ebf0ae4364ae86f2 Mon Sep 17 00:00:00 2001 From: Nikhil Ponnuru Date: Tue, 5 Aug 2025 15:18:49 +0530 Subject: [PATCH] feat: improve login tool behavior after authentication What it does: - Improves user messaging when login tool is called after authentication - Addresses MCP protocol limitations around tool discovery timing Why this decision: - MCP Limitation: Tools are discovered once at connection time, not dynamically - We can't truly hide the login tool, so we check if already logged in or not before using login tool --- mcp/mcp.go | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/mcp/mcp.go b/mcp/mcp.go index f4beb2f..ef49391 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -1,6 +1,7 @@ package mcp import ( + "context" "log/slog" "strings" @@ -96,9 +97,14 @@ func RegisterTools(srv *server.MCPServer, manager *kc.Manager, excludedTools str allTools := GetAllTools() filteredTools, registeredCount, excludedCount := filterTools(allTools, excludedSet) - // Register filtered tools + // Register filtered tools with session-aware login tool handling for _, tool := range filteredTools { - srv.AddTool(tool.Tool(), tool.Handler(manager)) + if tool.Tool().Name == "login" { + // Register login tool with session-aware handler that can reject calls + srv.AddTool(tool.Tool(), createSessionAwareLoginHandler(tool.Handler(manager), manager, logger)) + } else { + srv.AddTool(tool.Tool(), tool.Handler(manager)) + } } logger.Info("Tool registration complete", @@ -106,3 +112,21 @@ func RegisterTools(srv *server.MCPServer, manager *kc.Manager, excludedTools str "excluded", excludedCount, "total_available", len(allTools)) } + +// createSessionAwareLoginHandler creates a handler that rejects login tool calls when user is authenticated +func createSessionAwareLoginHandler(originalHandler server.ToolHandlerFunc, manager *kc.Manager, logger *slog.Logger) server.ToolHandlerFunc { + return func(ctx context.Context, request gomcp.CallToolRequest) (*gomcp.CallToolResult, error) { + mcpSession := server.ClientSessionFromContext(ctx) + sessionID := mcpSession.SessionID() + + // Check if user has valid session credentials + if _, err := manager.GetAuthenticatedClient(sessionID); err == nil { + // User is authenticated - reject the tool call + logger.Info("Login tool call rejected - user has valid session", "session_id", sessionID) + return gomcp.NewToolResultError("Tool unavailable: User is already authenticated. The login tool is disabled for authenticated sessions."), nil + } + + // User is not authenticated, proceed with original handler + return originalHandler(ctx, request) + } +}