Merge pull request #30 from nimishamehta5/rate-limit-session-id

Implement rate-limiting based on session ID

Author: JAORMX

README.md

@@ -15,6 +15,7 @@ MKP is a Model Context Protocol (MCP) server for Kubernetes that allows LLM-powe
 - Apply (create or update) clustered resources
 - Apply (create or update) namespaced resources
 - Generic and pluggable implementation using API Machinery's unstructured client
+- Built-in rate limiting for protection against excessive API calls
 
 ## Why MKP?
 
@@ -47,6 +48,7 @@ MKP offers several key advantages as a Model Context Protocol server for Kuberne
 ### Production-Ready Architecture
 - Designed for reliability and performance in production environments
 - Proper error handling and resource management
+- Built-in rate limiting to protect against excessive API calls
 - Testable design with comprehensive unit tests
 - Follows Kubernetes development best practices
 
@@ -291,6 +293,24 @@ By default, MKP operates in read-only mode, meaning it does not allow write oper
 ./build/mkp-server --kubeconfig=/path/to/kubeconfig --read-write=true
 ```
 
+### Rate Limiting
+
+MKP includes a built-in rate limiting mechanism to protect the server from excessive API calls, which is particularly important when used with AI agents. The rate limiter uses a token bucket algorithm and applies different limits based on the operation type:
+
+- Read operations (list_resources, get_resource): 120 requests per minute
+- Write operations (apply_resource, delete_resource): 30 requests per minute
+- Default for other operations: 60 requests per minute
+
+Rate limits are applied per client session, ensuring fair resource allocation across multiple clients. The rate limiting feature can be enabled or disabled via the command line flag:
+
+```bash
+# Run with rate limiting enabled (default)
+./build/mkp-server
+
+# Run with rate limiting disabled
+./build/mkp-server --enable-rate-limiting=false
+```
+
 ## Development
 
 ### Running tests

cmd/server/main.go

@@ -24,6 +24,9 @@ func main() {
 		"Whether to allow write operations on the cluster. When false, the server operates in read-only mode")
 	kubeconfigRefreshInterval := flag.Duration("kubeconfig-refresh-interval", 0,
 		"Interval to periodically re-read the kubeconfig (e.g., 5m for 5 minutes). If 0, no refresh will be performed")
+	enableRateLimiting := flag.Bool("enable-rate-limiting", true,
+		"Whether to enable rate limiting for tool calls. When false, no rate limiting will be applied")
+
 	flag.Parse()
 
 	// Create a context that can be cancelled
@@ -61,8 +64,9 @@ func main() {
 
 	// Create MCP server config
 	config := &mcp.Config{
-		ServeResources: *serveResources,
-		ReadWrite:      *readWrite,
+		ServeResources:     *serveResources,
+		ReadWrite:          *readWrite,
+		EnableRateLimiting: *enableRateLimiting,
 	}
 
 	// Create MCP server using the helper function
@@ -99,10 +103,17 @@ func main() {
 	shutdownCh := make(chan error, 1)
 	go func() {
 		log.Println("Initiating server shutdown...")
+
+		// Stop the SSE server
 		err := sseServer.Shutdown(shutdownCtx)
 		if err != nil {
 			log.Printf("Error during shutdown: %v", err)
 		}
+
+		// Stop the MCP server resources (including rate limiter)
+		log.Println("Stopping MCP server resources...")
+		mcp.StopServer()
+
 		shutdownCh <- err
 		close(shutdownCh)
 	}()

go.mod

@@ -32,6 +32,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/spf13/cast v1.7.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	golang.org/x/net v0.38.0 // indirect

pkg/mcp/apply_resource.go

@@ -9,6 +9,8 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+
+	"github.com/StacklokLabs/mkp/pkg/types"
 )
 
 // HandleApplyResource handles the apply_resource tool
@@ -31,7 +33,7 @@ func (m *Implementation) HandleApplyResource(ctx context.Context, request mcp.Ca
 	if resource == "" {
 		return mcp.NewToolResultError("resource is required"), nil
 	}
-	if resourceType == ResourceTypeNamespaced && namespace == "" {
+	if resourceType == types.ResourceTypeNamespaced && namespace == "" {
 		return mcp.NewToolResultError("namespace is required for namespaced resources"), nil
 	}
 	if manifestMap == nil {
@@ -52,9 +54,9 @@ func (m *Implementation) HandleApplyResource(ctx context.Context, request mcp.Ca
 	var result *unstructured.Unstructured
 	var err error
 	switch resourceType {
-	case ResourceTypeClustered:
+	case types.ResourceTypeClustered:
 		result, err = m.k8sClient.ApplyClusteredResource(ctx, gvr, obj)
-	case ResourceTypeNamespaced:
+	case types.ResourceTypeNamespaced:
 		result, err = m.k8sClient.ApplyNamespacedResource(ctx, gvr, namespace, obj)
 	default:
 		return mcp.NewToolResultError(fmt.Sprintf("Invalid resource_type: %s", resourceType)), nil

pkg/mcp/apply_resource_test.go

@@ -13,6 +13,7 @@ import (
 	ktesting "k8s.io/client-go/testing"
 
 	"github.com/StacklokLabs/mkp/pkg/k8s"
+	"github.com/StacklokLabs/mkp/pkg/types"
 )
 
 func TestHandleApplyResourceClusteredSuccess(t *testing.T) {
@@ -59,9 +60,9 @@ func TestHandleApplyResourceClusteredSuccess(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = ApplyResourceToolName
+	request.Params.Name = types.ApplyResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": "clustered",
+		"resource_type": types.ResourceTypeClustered,
 		"group":         "rbac.authorization.k8s.io",
 		"version":       "v1",
 		"resource":      "clusterroles",
@@ -146,9 +147,9 @@ func TestHandleApplyResourceNamespacedSuccess(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = ApplyResourceToolName
+	request.Params.Name = types.ApplyResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": ResourceTypeNamespaced,
+		"resource_type": types.ResourceTypeNamespaced,
 		"group":         "",
 		"version":       "v1",
 		"resource":      "services",
@@ -216,7 +217,7 @@ func TestHandleApplyResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing version",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"resource":      "deployments",
 				"manifest":      map[string]interface{}{},
@@ -226,7 +227,7 @@ func TestHandleApplyResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing resource",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"version":       "v1",
 				"manifest":      map[string]interface{}{},
@@ -236,7 +237,7 @@ func TestHandleApplyResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing namespace for namespaced resource",
 			arguments: map[string]interface{}{
-				"resource_type": ResourceTypeNamespaced,
+				"resource_type": types.ResourceTypeNamespaced,
 				"group":         "apps",
 				"version":       "v1",
 				"resource":      "deployments",
@@ -247,7 +248,7 @@ func TestHandleApplyResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing manifest",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"version":       "v1",
 				"resource":      "deployments",
@@ -260,7 +261,7 @@ func TestHandleApplyResourceMissingParameters(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			// Create a test request
 			request := mcp.CallToolRequest{}
-			request.Params.Name = ApplyResourceToolName
+			request.Params.Name = types.ApplyResourceToolName
 			request.Params.Arguments = tc.arguments
 
 			// Test HandleApplyResource
@@ -293,7 +294,7 @@ func TestHandleApplyResourceInvalidResourceType(t *testing.T) {
 
 	// Create a test request with invalid resource_type
 	request := mcp.CallToolRequest{}
-	request.Params.Name = ApplyResourceToolName
+	request.Params.Name = types.ApplyResourceToolName
 	request.Params.Arguments = map[string]interface{}{
 		"resource_type": "invalid",
 		"group":         "apps",
@@ -347,9 +348,9 @@ func TestHandleApplyResourceApplyError(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = ApplyResourceToolName
+	request.Params.Name = types.ApplyResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": "clustered",
+		"resource_type": types.ResourceTypeClustered,
 		"group":         "rbac.authorization.k8s.io",
 		"version":       "v1",
 		"resource":      "clusterroles",

pkg/mcp/delete_resource.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/mark3labs/mcp-go/mcp"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+
+	"github.com/StacklokLabs/mkp/pkg/types"
 )
 
 // HandleDeleteResource handles the delete_resource tool
@@ -32,13 +34,13 @@ func (m *Implementation) HandleDeleteResource(ctx context.Context, request mcp.C
 	if name == "" {
 		return mcp.NewToolResultError("name is required"), nil
 	}
-	if resourceType == ResourceTypeNamespaced && namespace == "" {
+	if resourceType == types.ResourceTypeNamespaced && namespace == "" {
 		return mcp.NewToolResultError("namespace is required for namespaced resources"), nil
 	}
 
 	// Create GVR
 	// Validate resource_type
-	if resourceType != ResourceTypeClustered && resourceType != ResourceTypeNamespaced {
+	if resourceType != types.ResourceTypeClustered && resourceType != types.ResourceTypeNamespaced {
 		return mcp.NewToolResultError("Invalid resource_type: " + resourceType), nil
 	}
 
@@ -51,9 +53,9 @@ func (m *Implementation) HandleDeleteResource(ctx context.Context, request mcp.C
 	// Delete resource
 	var err error
 	switch resourceType {
-	case ResourceTypeClustered:
+	case types.ResourceTypeClustered:
 		err = m.k8sClient.DeleteClusteredResource(ctx, gvr, name)
-	case ResourceTypeNamespaced:
+	case types.ResourceTypeNamespaced:
 		err = m.k8sClient.DeleteNamespacedResource(ctx, gvr, namespace, name)
 	default:
 		return mcp.NewToolResultError(fmt.Sprintf("Invalid resource_type: %s", resourceType)), nil
@@ -68,7 +70,7 @@ func (m *Implementation) HandleDeleteResource(ctx context.Context, request mcp.C
 
 // NewDeleteResourceTool creates a new delete_resource tool
 func NewDeleteResourceTool() mcp.Tool {
-	return mcp.NewTool("delete_resource",
+	return mcp.NewTool(types.DeleteResourceToolName,
 		mcp.WithDescription("Delete a Kubernetes resource"),
 		mcp.WithString("resource_type",
 			mcp.Description("Type of resource to delete (clustered or namespaced)"),

pkg/mcp/delete_resource_test.go

@@ -12,6 +12,7 @@ import (
 	ktesting "k8s.io/client-go/testing"
 
 	"github.com/StacklokLabs/mkp/pkg/k8s"
+	"github.com/StacklokLabs/mkp/pkg/types"
 )
 
 func TestHandleDeleteResourceClusteredSuccess(t *testing.T) {
@@ -35,9 +36,9 @@ func TestHandleDeleteResourceClusteredSuccess(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = DeleteResourceToolName
+	request.Params.Name = types.DeleteResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": "clustered",
+		"resource_type": types.ResourceTypeClustered,
 		"group":         "rbac.authorization.k8s.io",
 		"version":       "v1",
 		"resource":      "clusterroles",
@@ -84,9 +85,9 @@ func TestHandleDeleteResourceNamespacedSuccess(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = DeleteResourceToolName
+	request.Params.Name = types.DeleteResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": ResourceTypeNamespaced,
+		"resource_type": types.ResourceTypeNamespaced,
 		"group":         "",
 		"version":       "v1",
 		"resource":      "services",
@@ -139,7 +140,7 @@ func TestHandleDeleteResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing version",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"resource":      "deployments",
 				"name":          "test-deployment",
@@ -149,7 +150,7 @@ func TestHandleDeleteResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing resource",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"version":       "v1",
 				"name":          "test-deployment",
@@ -159,7 +160,7 @@ func TestHandleDeleteResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing name",
 			arguments: map[string]interface{}{
-				"resource_type": "clustered",
+				"resource_type": types.ResourceTypeClustered,
 				"group":         "apps",
 				"version":       "v1",
 				"resource":      "deployments",
@@ -169,7 +170,7 @@ func TestHandleDeleteResourceMissingParameters(t *testing.T) {
 		{
 			name: "Missing namespace for namespaced resource",
 			arguments: map[string]interface{}{
-				"resource_type": ResourceTypeNamespaced,
+				"resource_type": types.ResourceTypeNamespaced,
 				"group":         "apps",
 				"version":       "v1",
 				"resource":      "deployments",
@@ -183,7 +184,7 @@ func TestHandleDeleteResourceMissingParameters(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			// Create a test request
 			request := mcp.CallToolRequest{}
-			request.Params.Name = DeleteResourceToolName
+			request.Params.Name = types.DeleteResourceToolName
 			request.Params.Arguments = tc.arguments
 
 			// Test HandleDeleteResource
@@ -216,7 +217,7 @@ func TestHandleDeleteResourceInvalidResourceType(t *testing.T) {
 
 	// Create a test request with invalid resource_type
 	request := mcp.CallToolRequest{}
-	request.Params.Name = DeleteResourceToolName
+	request.Params.Name = types.DeleteResourceToolName
 	request.Params.Arguments = map[string]interface{}{
 		"resource_type": "invalid",
 		"group":         "apps",
@@ -265,9 +266,9 @@ func TestHandleDeleteResourceDeleteError(t *testing.T) {
 
 	// Create a test request
 	request := mcp.CallToolRequest{}
-	request.Params.Name = DeleteResourceToolName
+	request.Params.Name = types.DeleteResourceToolName
 	request.Params.Arguments = map[string]interface{}{
-		"resource_type": "clustered",
+		"resource_type": types.ResourceTypeClustered,
 		"group":         "rbac.authorization.k8s.io",
 		"version":       "v1",
 		"resource":      "clusterroles",

pkg/mcp/get_resource.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/mark3labs/mcp-go/mcp"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+
+	"github.com/StacklokLabs/mkp/pkg/types"
 )
 
 // HandleGetResource handles the get_resource tool
@@ -50,13 +52,13 @@ func (m *Implementation) HandleGetResource(ctx context.Context, request mcp.Call
 	if name == "" {
 		return mcp.NewToolResultError("name is required"), nil
 	}
-	if resourceType == ResourceTypeNamespaced && namespace == "" {
+	if resourceType == types.ResourceTypeNamespaced && namespace == "" {
 		return mcp.NewToolResultError("namespace is required for namespaced resources"), nil
 	}
 
 	// Create GVR
 	// Validate resource_type
-	if resourceType != "clustered" && resourceType != ResourceTypeNamespaced {
+	if resourceType != types.ResourceTypeClustered && resourceType != types.ResourceTypeNamespaced {
 		return mcp.NewToolResultError("Invalid resource_type: " + resourceType), nil
 	}
 
@@ -83,7 +85,7 @@ func (m *Implementation) HandleGetResource(ctx context.Context, request mcp.Call
 
 // NewGetResourceTool creates a new get_resource tool
 func NewGetResourceTool() mcp.Tool {
-	return mcp.NewTool("get_resource",
+	return mcp.NewTool(types.GetResourceToolName,
 		mcp.WithDescription("Get a Kubernetes resource or its subresource"),
 		mcp.WithString("resource_type",
 			mcp.Description("Type of resource to get (clustered or namespaced)"),