fix(tok): fix unused variable, limits, mutexes

chunker.go

@@ -67,14 +67,18 @@ var (
 var (
 	customLanguagePatterns   = map[string][]string{} // lang -> boundary regexes
 	customLanguagePatternsMu sync.RWMutex
+	compiledCustomBoundary   sync.Map // lang -> *regexp.Regexp
 )
 
 // RegisterLanguagePatterns registers custom boundary patterns for a language.
 // These override the built-in patterns when looking up boundaries.
 func RegisterLanguagePatterns(lang string, patterns []string) {
 	customLanguagePatternsMu.Lock()
-	defer customLanguagePatternsMu.Unlock()
 	customLanguagePatterns[strings.ToLower(lang)] = patterns
+	customLanguagePatternsMu.Unlock()
+
+	// Invalidate cached compiled regex for this language
+	compiledCustomBoundary.Delete(strings.ToLower(lang))
 }
 
 // GetLanguagePatterns returns custom patterns if registered, else built-in pattern strings.
@@ -132,11 +136,12 @@ func splitBySepLevel(content string, maxTokens int, level int) []string {
 	}
 
 	// Greedily merge pieces to stay within maxTokens
+	// Use fast estimate for intermediate decisions to avoid O(n²) BPE cost.
 	var result []string
 	current := pieces[0]
 	for i := 1; i < len(pieces); i++ {
 		combined := current + pieces[i]
-		if EstimateTokensPrecise(combined) <= maxTokens {
+		if EstimateTokensFast(combined) <= maxTokens {
 			current = combined
 		} else {
 			result = append(result, current)
@@ -367,15 +372,21 @@ func greedyMerging(chunks []rawChunk, opts ChunkOptions) []rawChunk {
 				continue
 			}
 
-			curTokens := EstimateTokensPrecise(result[i].content)
-			nextTokens := EstimateTokensPrecise(result[i+1].content)
 			combined := result[i].content + "\n" + result[i+1].content
-			combinedTokens := EstimateTokensPrecise(combined)
+			combinedTokens := EstimateTokensFast(combined)
+			if combinedTokens > opts.MaxTokens {
+				next = append(next, result[i])
+				i++
+				continue
+			}
+
+			curTokens := EstimateTokensFast(result[i].content)
+			nextTokens := EstimateTokensFast(result[i+1].content)
 
 			costSeparate := chunkCost(curTokens, opts) + chunkCost(nextTokens, opts)
 			costMerged := chunkCost(combinedTokens, opts)
 
-			if costMerged < costSeparate && combinedTokens <= opts.MaxTokens {
+			if costMerged < costSeparate {
 				merged := rawChunk{
 					content:   combined,
 					startLine: result[i].startLine,
@@ -733,10 +744,15 @@ func boundaryForLang(lang string) *regexp.Regexp {
 	customLanguagePatternsMu.RUnlock()
 
 	if ok && len(patterns) > 0 {
-		// Combine custom patterns into a single alternation regex
+		if cached, found := compiledCustomBoundary.Load(lang); found {
+			if re, ok := cached.(*regexp.Regexp); ok {
+				return re
+			}
+		}
 		combined := strings.Join(patterns, "|")
 		re, err := regexp.Compile(combined)
 		if err == nil {
+			compiledCustomBoundary.Store(lang, re)
 			return re
 		}
 	}

internal/cache/query_cache.go

@@ -72,6 +72,8 @@ func NewQueryCache(dbPath string) (*QueryCache, error) {
 		return nil, fmt.Errorf("open cache database: %w", err)
 	}
 
+	db.SetMaxOpenConns(1)
+
 	// Set pragmas for performance
 	if _, err := db.ExecContext(context.Background(), `
 		PRAGMA journal_mode = WAL;

mcp/server.go

@@ -19,11 +19,15 @@ package mcp
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"sync"
 
 	tok "github.com/GrayCodeAI/tok"
 )
 
+const maxInputSize = 16 << 20 // 16 MiB
+
 type ToolHandler func(ctx context.Context, params map[string]interface{}) (interface{}, error)
 
 type ToolDef struct {
@@ -35,6 +39,7 @@ type ToolDef struct {
 type MCPServer struct {
 	name  string
 	tools map[string]toolEntry
+	mu    sync.RWMutex
 }
 
 type toolEntry struct {
@@ -51,18 +56,32 @@ func NewServer(name string) *MCPServer {
 // RegisterTool adds a tool to the server. Re-registering an existing name
 // overwrites the previous definition and handler.
 func (s *MCPServer) RegisterTool(name, description string, schema map[string]interface{}, handler ToolHandler) {
+	s.mu.Lock()
 	s.tools[name] = toolEntry{def: ToolDef{Name: name, Description: description, InputSchema: schema}, handler: handler}
+	s.mu.Unlock()
 }
 
 // ListTools returns the registered tool definitions. Order is not stable.
 func (s *MCPServer) ListTools() []ToolDef {
+	s.mu.RLock()
 	defs := make([]ToolDef, 0, len(s.tools))
 	for _, e := range s.tools {
 		defs = append(defs, e.def)
 	}
+	s.mu.RUnlock()
 	return defs
 }
 
+var ErrInputTooLarge = errors.New("input exceeds maximum size (16 MiB)")
+
+// checkTextSize validates that a text parameter does not exceed maxInputSize.
+func checkTextSize(text string) error {
+	if len(text) > maxInputSize {
+		return ErrInputTooLarge
+	}
+	return nil
+}
+
 // HandleRequest dispatches a JSON-RPC-style request to the registered tool.
 // Supported methods: "tools/list" (no params) and "tools/call" (params must
 // include "name" and optional "arguments" map).
@@ -75,7 +94,9 @@ func (s *MCPServer) HandleRequest(ctx context.Context, method string, params map
 		if name == "" {
 			return nil, fmt.Errorf("tools/call: missing tool name")
 		}
+		s.mu.RLock()
 		entry, ok := s.tools[name]
+		s.mu.RUnlock()
 		if !ok {
 			return nil, fmt.Errorf("tools/call: unknown tool %q", name)
 		}
@@ -162,6 +183,9 @@ func countTokensHandler(_ context.Context, p map[string]interface{}) (interface{
 	if text == "" {
 		return nil, fmt.Errorf("count_tokens: text required")
 	}
+	if err := checkTextSize(text); err != nil {
+		return nil, err
+	}
 	model, _ := p["model"].(string)
 	if model == "" || model == "heuristic" {
 		return map[string]interface{}{
@@ -222,6 +246,9 @@ func compressTextHandler(_ context.Context, p map[string]interface{}) (interface
 	if text == "" {
 		return nil, fmt.Errorf("compress_text: text required")
 	}
+	if err := checkTextSize(text); err != nil {
+		return nil, err
+	}
 	mode, _ := p["mode"].(string)
 	opts := buildCompressOptions(mode, p)
 	out, stats := tok.Compress(text, opts...)
@@ -246,6 +273,9 @@ func redactSecretsHandler(_ context.Context, p map[string]interface{}) (interfac
 	if text == "" {
 		return nil, fmt.Errorf("redact_secrets: text required")
 	}
+	if err := checkTextSize(text); err != nil {
+		return nil, err
+	}
 	entropy, hasEntropy := numberFromParams(p, "entropyThreshold")
 	det := tok.NewSecretDetector()
 	var matches []tok.SecretMatch