limits on number of files and file size when loading or updating cont…

…ext, closes #196

app/cli/cmd/tell.go

@@ -16,6 +16,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// test
 const defaultEditor = "vim"
 
 // const defaultEditor = "nano"

app/cli/lib/context_load.go

@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"encoding/base64"
 	"fmt"
-
 	"io"
 	"os"
 	"plandex/api"
@@ -215,6 +214,11 @@ func MustLoadContext(resources []string, params *types.LoadContextParams) {
 				onErr(fmt.Errorf("failed to parse input paths: %v", err))
 			}
 
+			// Add this check for the number of files
+			if len(flattenedPaths) > shared.MaxContextCount {
+				onErr(fmt.Errorf("too many files to load (found %d, limit is %d)", len(flattenedPaths), shared.MaxContextCount))
+			}
+
 			if !params.ForceSkipIgnore {
 				var filteredPaths []string
 				for _, path := range flattenedPaths {
@@ -249,6 +253,17 @@ func MustLoadContext(resources []string, params *types.LoadContextParams) {
 
 				numRoutines++
 				go func(path string) {
+					// File size check
+					fileInfo, err := os.Stat(path)
+					if err != nil {
+						errCh <- fmt.Errorf("failed to get file info for %s: %v", path, err)
+						return
+					}
+
+					if fileInfo.Size() > shared.MaxContextBodySize {
+						errCh <- fmt.Errorf("file %s exceeds size limit (size %.2f MB, limit %d MB)", path, float64(fileInfo.Size())/1024/1024, int(shared.MaxContextBodySize)/1024/1024)
+						return
+					}
 
 					fileContent, err := os.ReadFile(path)
 					if err != nil {

app/cli/lib/context_update.go

@@ -380,6 +380,28 @@ func checkOutdatedAndMaybeUpdateContext(doUpdate bool, maybeContexts []*shared.C
 		return nil, fmt.Errorf("failed to check context outdated: %v", errs)
 	}
 
+	var totalContextCount int
+	var totalBodySize int64
+
+	for _, context := range contexts {
+		totalContextCount++
+		totalBodySize += int64(len(context.Body))
+	}
+
+	for _, context := range updatedContexts {
+		if req[context.Id] != nil {
+			totalBodySize += int64(len(req[context.Id].Body)) - int64(len(context.Body))
+		}
+	}
+
+	if totalContextCount > shared.MaxContextCount {
+		return nil, fmt.Errorf("too many contexts to update (found %d, limit is %d)", totalContextCount, shared.MaxContextCount)
+	}
+
+	if totalBodySize > shared.MaxContextBodySize {
+		return nil, fmt.Errorf("total context body size exceeds limit (size %.2f MB, limit %d MB)", float64(totalBodySize)/1024/1024, int(shared.MaxContextBodySize)/1024/1024)
+	}
+
 	var msg string
 	var hasConflicts bool
 

app/server/db/context_helpers.go

@@ -371,6 +371,31 @@ func UpdateContexts(params UpdateContextsParams) (*shared.UpdateContextResponse,
 		contextsById = params.ContextsById
 	}
 
+	var totalContextCount int
+	var totalBodySize int64
+
+	for _, context := range contextsById {
+		totalContextCount++
+		totalBodySize += int64(len(context.Body))
+	}
+
+	for id, params := range *req {
+		if context, ok := contextsById[id]; ok {
+			totalBodySize += int64(len(params.Body)) - int64(len(context.Body))
+		} else {
+			totalContextCount++
+			totalBodySize += int64(len(params.Body))
+		}
+	}
+
+	if totalContextCount > shared.MaxContextCount {
+		return nil, fmt.Errorf("too many contexts to update (found %d, limit is %d)", totalContextCount, shared.MaxContextCount)
+	}
+
+	if totalBodySize > shared.MaxContextBodySize {
+		return nil, fmt.Errorf("total context body size exceeds limit (size %.2f MB, limit %d MB)", float64(totalBodySize)/1024/1024, int(shared.MaxContextBodySize)/1024/1024)
+	}
+
 	var updatedContexts []*shared.Context
 
 	numFiles := 0

app/server/handlers/context_helper.go

@@ -15,6 +15,24 @@ import (
 )
 
 func loadContexts(w http.ResponseWriter, r *http.Request, auth *types.ServerAuth, loadReq *shared.LoadContextRequest, plan *db.Plan, branchName string) (*shared.LoadContextResponse, []*db.Context) {
+	// check file count and size limits
+	totalFiles := 0
+	for _, context := range *loadReq {
+		totalFiles++
+		if totalFiles > shared.MaxContextCount {
+			log.Printf("Error: Too many contexts to load (found %d, limit is %d)\n", totalFiles, shared.MaxContextCount)
+			http.Error(w, fmt.Sprintf("Too many contexts to load (found %d, limit is %d)", totalFiles, shared.MaxContextCount), http.StatusBadRequest)
+			return nil, nil
+		}
+
+		fileSize := int64(len(context.Body))
+		if fileSize > shared.MaxContextBodySize {
+			log.Printf("Error: Context %s exceeds size limit (size %.2f MB, limit %d MB)\n", context.Name, float64(fileSize)/1024/1024, int(shared.MaxContextBodySize)/1024/1024)
+			http.Error(w, fmt.Sprintf("Context %s exceeds size limit (size %.2f MB, limit %d MB)", context.Name, float64(fileSize)/1024/1024, int(shared.MaxContextBodySize)/1024/1024), http.StatusBadRequest)
+			return nil, nil
+		}
+	}
+
 	var err error
 	var settings *shared.PlanSettings
 	var client *openai.Client

app/server/setup/setup.go

@@ -67,27 +67,31 @@ func StartServer(handler http.Handler) {
 		externalPort = "8080"
 	}
 
+	// Apply the maxBytesMiddleware to limit request size to 100 MB
+	handler = maxBytesMiddleware(handler, 100<<20) // 100 MB limit
+
 	// Enable CORS based on environment
-	var corsHandler http.Handler
 	if os.Getenv("GOENV") == "development" {
-		corsHandler = cors.New(cors.Options{
+		handler = cors.New(cors.Options{
 			AllowedOrigins:   []string{"*"},
-			AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE"},
+			AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
 			AllowedHeaders:   []string{"Content-Type", "Authorization"},
 			AllowCredentials: true,
 		}).Handler(handler)
 	} else {
-		corsHandler = cors.New(cors.Options{
+		handler = cors.New(cors.Options{
 			AllowedOrigins:   []string{fmt.Sprintf("https://%s.plandex.ai", os.Getenv("APP_SUBDOMAIN"))},
-			AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE"},
+			AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
 			AllowedHeaders:   []string{"Content-Type", "Authorization"},
 			AllowCredentials: true,
 		}).Handler(handler)
 	}
 
 	server := &http.Server{
-		Addr:    ":" + externalPort,
-		Handler: corsHandler,
+		Addr:              ":" + externalPort,
+		Handler:           handler,
+		MaxHeaderBytes:    1 << 20, // 1 MB
+		ReadHeaderTimeout: 5 * time.Second,
 	}
 
 	go func() {
@@ -148,3 +152,10 @@ func waitForActivePlans() chan struct{} {
 	}()
 	return done
 }
+
+func maxBytesMiddleware(next http.Handler, maxBytes int64) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
+		next.ServeHTTP(w, r)
+	})
+}

app/shared/context.go

@@ -9,6 +9,11 @@ import (
 	"github.com/olekukonko/tablewriter"
 )
 
+const (
+	MaxContextBodySize = 10 * 1024 * 1024 // 10MB
+	MaxContextCount    = 500
+)
+
 type ContextUpdateResult struct {
 	UpdatedContexts []*Context
 	TokenDiffsById  map[string]int