diff --git a/chunker.go b/chunker.go index ed3ece29b..5e14b1d8f 100644 --- a/chunker.go +++ b/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 } } diff --git a/internal/cache/query_cache.go b/internal/cache/query_cache.go index 69ad4d128..82a8d309f 100644 --- a/internal/cache/query_cache.go +++ b/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; diff --git a/mcp/server.go b/mcp/server.go index 27aab049f..8a9cdb015 100644 --- a/mcp/server.go +++ b/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