@getmikk/ai-context 1.9.0 → 2.0.0

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@getmikk/ai-context",
-    "version": "1.9.0",
+    "version": "2.0.0",
     "license": "Apache-2.0",
     "repository": {
         "type": "git",
@@ -21,8 +21,8 @@
         "dev": "tsc --watch"
     },
     "dependencies": {
-        "@getmikk/core": "^1.9.0",
-        "@getmikk/intent-engine": "^1.9.0"
+        "@getmikk/core": "^2.0.0",
+        "@getmikk/intent-engine": "^2.0.0"
     },
     "devDependencies": {
         "typescript": "^5.7.0",

package/src/context-builder.ts

@@ -373,16 +373,24 @@ export class ContextBuilder {
 
         // ── Step 5: Fill token budget ──────────────────────────────────────
         let selected: MikkLockFunction[] = []
+        
+        // Pre-calculate baseline overhead (context files, routes, constraints)
         let usedTokens = 0
+        const routesStr = (!strictMode && this.lock.routes) ? JSON.stringify(this.lock.routes) : ''
+        const ctxStr = (!strictMode && this.lock.contextFiles) 
+            ? this.lock.contextFiles.map(cf => readContextFile(cf.path, query.projectRoot).slice(0, 2000)).join('\n')
+            : ''
+        usedTokens += estimateTokens(routesStr + ctxStr + JSON.stringify(this.contract.declared.constraints))
 
         for (const { fn, score } of scored) {
             if (score <= 0 && seeds.length > 0) break // Nothing relevant left
             if (selected.length >= (query.maxFunctions ?? 80)) break
 
             const snippet = this.buildFunctionSnippet(fn, query)
-            const tokens = estimateTokens(snippet)
+            // Multiply tokens by 2.2 to account for it being in both JSON and text prompt, plus JSON framing
+            const tokens = estimateTokens(snippet) * 2.2
 
-            if (usedTokens + tokens > tokenBudget) continue  // skip, try smaller ones later
+            if (usedTokens + tokens > tokenBudget && selected.length > 0) continue  // skip, try smaller ones later
             selected.push(fn)
             usedTokens += tokens
         }

package/src/token-counter.ts

@@ -1,224 +1,157 @@
 /**
- * Improved Token Counter
- * 
- * Provides more accurate token counting than the simple length/4 approximation.
- * Uses a GPT-4 compatible tokenizer approximation for better budget management.
+ * Token Counter — accurate, fast token estimation for context budget management.
+ *
+ * Design:
+ *  - `countTokens(text)` — accurate, linear-scan, O(n)
+ *  - `countTokensFast(text)` — single-pass heuristic, O(n) for hot paths
+ *  - `estimateFileTokens(content, path)` — file-type-aware wrapper
+ *  - `TokenBudget` — budget manager with truncation
+ *
+ * The previous implementation used a character-position Set to track processed
+ * ranges across multiple regex scans — O(n²) per call on large files.
+ * Replaced with a single linear scan that categorises characters without
+ * per-character Set lookups.
  */
 
-// Character-based token approximation (more accurate than simple division)
-const CHARS_PER_TOKEN = 3.8 // Average for GPT-4 tokenizer
-const MIN_CHARS_PER_TOKEN = 2.0   // For dense code
-const MAX_CHARS_PER_TOKEN = 6.0   // For sparse text
-
-// Special token patterns that affect tokenization
-const TOKEN_PATTERNS = {
-    // Common programming patterns that typically tokenize as single tokens
-    SINGLE_TOKEN_PATTERNS: [
-        /\b(if|else|for|while|function|return|const|let|var|class|import|export)\b/g,
-        /\b(true|false|null|undefined)\b/g,
-        /\b(async|await|try|catch|throw|new|this)\b/g,
-        // Operators and punctuation
-        /[+\-*\/=<>!&|]+/g,
-        /[{}()\[\];,\.]/g,
-        // Common function names
-        /\b(console\.log|console\.error|console\.warn)\b/g,
-        /\b(Math\.(floor|ceil|round|max|min))\b/g,
-    ],
-    
-    // Patterns that typically increase token count
-    HIGH_TOKEN_PATTERNS: [
-        // String literals (each character ~0.25 tokens)
-        /'[^']*'/g,
-        /"[^"]*"/g,
-        /`[^`]*`/g,
-        // Numbers (digits ~0.5 tokens each)
-        /\b\d+\.?\d*\b/g,
-        // Long identifiers (split into multiple tokens)
-        /\b[a-z][a-zA-Z0-9]{8,}\b/g,
-    ]
-}
+const CHARS_PER_TOKEN     = 3.8  // GPT-4 average
+const MIN_CHARS_PER_TOKEN = 2.0  // Dense code
+const MAX_CHARS_PER_TOKEN = 6.0  // Sparse natural language
 
 /**
- * Count tokens with improved accuracy using position-based pattern matching
+ * Count tokens with reasonable accuracy — O(n) single linear scan.
+ *
+ * Classifies runs of characters into:
+ *   - whitespace: free (separators, not tokens)
+ *   - string literals: ~4 chars/token
+ *   - digit runs: ~2 chars/token (numbers tokenise finely)
+ *   - identifiers/keywords: short → 1 token, long → ~3.5 chars/token
+ *   - operators/punctuation: 1 char = 1 token
  */
 export function countTokens(text: string): number {
-    if (!text || text.length === 0) return 0
-    
-    let tokenCount = 0
-    const processedPositions = new Set<number>() // Track positions to avoid double-counting
-    
-    // Count single-token patterns with position tracking
-    for (const pattern of TOKEN_PATTERNS.SINGLE_TOKEN_PATTERNS) {
-        for (const match of text.matchAll(pattern)) {
-            const start = match.index!
-            const end = start + match[0].length
-            
-            // Check if this range overlaps with already processed ranges
-            let overlaps = false
-            for (let i = start; i < end; i++) {
-                if (processedPositions.has(i)) {
-                    overlaps = true
-                    break
-                }
-            }
-            
-            if (!overlaps) {
-                tokenCount += 1
-                // Mark positions as processed
-                for (let i = start; i < end; i++) {
-                    processedPositions.add(i)
-                }
-            }
+    if (!text) return 0
+
+    let tokens = 0
+    let i = 0
+    const n = text.length
+
+    while (i < n) {
+        const ch = text[i]
+
+        // Whitespace — boundary only, no token cost
+        if (ch === ' ' || ch === '\t' || ch === '\n' || ch === '\r') {
+            i++
+            continue
         }
-    }
-    
-    // Count high-token patterns (strings, numbers, long identifiers)
-    for (const pattern of TOKEN_PATTERNS.HIGH_TOKEN_PATTERNS) {
-        for (const match of text.matchAll(pattern)) {
-            const start = match.index!
-            const end = start + match[0].length
-            
-            // Check for overlaps
-            let overlaps = false
-            for (let i = start; i < end; i++) {
-                if (processedPositions.has(i)) {
-                    overlaps = true
-                    break
-                }
+
+        // String literals — scan to closing quote
+        if (ch === '"' || ch === "'" || ch === '`') {
+            const q = ch
+            let len = 1
+            i++
+            while (i < n) {
+                if (text[i] === '\\') { i += 2; len += 2; continue }
+                if (text[i] === q) { i++; len++; break }
+                i++; len++
             }
-            
-            if (!overlaps) {
-                let tokensToAdd = 0
-                if (match[0].startsWith('\'') || match[0].startsWith('"') || match[0].startsWith('`')) {
-                    // String literal: roughly 1 token per 4 characters
-                    tokensToAdd = Math.ceil(match[0].length / 4)
-                } else if (/^\d/.test(match[0])) {
-                    // Number: roughly 1 token per 2 digits
-                    tokensToAdd = Math.ceil(match[0].length / 2)
-                } else {
-                    // Long identifier: roughly 1 token per 6 characters
-                    tokensToAdd = Math.ceil(match[0].length / 6)
-                }
-                
-                tokenCount += tokensToAdd
-                // Mark positions as processed
-                for (let i = start; i < end; i++) {
-                    processedPositions.add(i)
-                }
+            tokens += Math.max(1, Math.ceil(len / 4))
+            continue
+        }
+
+        // Digit runs — token-heavy
+        if (ch >= '0' && ch <= '9') {
+            let len = 0
+            while (i < n && ((text[i] >= '0' && text[i] <= '9') || text[i] === '.')) {
+                i++; len++
             }
+            tokens += Math.max(1, Math.ceil(len / 2))
+            continue
         }
-    }
-    
-    // Count remaining characters (general text)
-    const remainingText = Array.from(text.split(''))
-        .map((char, index) => processedPositions.has(index) ? '' : char)
-        .join('')
-    
-    if (remainingText.length > 0) {
-        // Use variable rate based on character density
-        const avgWordLength = remainingText.split(/\s+/).reduce((sum, word) => sum + word.length, 0) / Math.max(remainingText.split(/\s+/).length, 1)
-        
-        let charsPerToken = CHARS_PER_TOKEN
-        if (avgWordLength < 4) {
-            charsPerToken = MIN_CHARS_PER_TOKEN // Dense code
-        } else if (avgWordLength > 8) {
-            charsPerToken = MAX_CHARS_PER_TOKEN // Sparse text
+
+        // Identifier / keyword runs
+        if ((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || ch === '_' || ch === '$') {
+            let len = 0
+            while (
+                i < n &&
+                ((text[i] >= 'a' && text[i] <= 'z') || (text[i] >= 'A' && text[i] <= 'Z') ||
+                 (text[i] >= '0' && text[i] <= '9') || text[i] === '_' || text[i] === '$')
+            ) { i++; len++ }
+            tokens += len <= 6 ? 1 : Math.ceil(len / 3.5)
+            continue
         }
-        
-        tokenCount += Math.ceil(remainingText.length / charsPerToken)
+
+        // Operators, punctuation, brackets — 1 char per token
+        tokens++
+        i++
     }
-    
-    // Apply bounds checking for sanity
+
     const minEstimate = Math.ceil(text.length / MAX_CHARS_PER_TOKEN)
     const maxEstimate = Math.ceil(text.length / MIN_CHARS_PER_TOKEN)
-    
-    return Math.max(minEstimate, Math.min(maxEstimate, tokenCount))
+    return Math.max(minEstimate, Math.min(maxEstimate, tokens))
 }
 
 /**
- * Fast token count for quick estimates (still more accurate than length/4)
+ * Fast O(n) single-pass heuristic for hot paths (context builder scoring loops).
  */
 export function countTokensFast(text: string): number {
-    if (!text || text.length === 0) return 0
-    
-    // Quick heuristic based on character patterns
-    const codeDensity = (text.match(/[a-zA-Z0-9]/g) || []).length / text.length
-    const stringRatio = (text.match(/['"`]/g) || []).length / text.length
-    
-    // Adjust chars per token based on content type
-    let charsPerToken = CHARS_PER_TOKEN
-    if (codeDensity > 0.7) {
-        charsPerToken = 3.2 // Dense code
-    } else if (stringRatio > 0.2) {
-        charsPerToken = 4.5 // String-heavy
-    } else if (codeDensity < 0.3) {
-        charsPerToken = 5.0 // Sparse text/comments
+    if (!text) return 0
+
+    let alphaNum = 0, punct = 0
+    for (let i = 0; i < text.length; i++) {
+        const c = text.charCodeAt(i)
+        if ((c >= 65 && c <= 90) || (c >= 97 && c <= 122) || (c >= 48 && c <= 57)) {
+            alphaNum++
+        } else if (c !== 32 && c !== 9 && c !== 10 && c !== 13) {
+            punct++
+        }
     }
-    
-    return Math.ceil(text.length / charsPerToken)
+
+    const nonWs = alphaNum + punct
+    if (nonWs === 0) return 0
+
+    const punctRatio = nonWs > 0 ? punct / nonWs : 0
+    const charsPerToken = punctRatio > 0.3 ? 2.8 : CHARS_PER_TOKEN
+    return Math.max(1, Math.ceil(text.length / charsPerToken))
 }
 
 /**
- * Estimate tokens for a file with content type awareness
+ * Estimate tokens for a file with content-type awareness.
  */
 export function estimateFileTokens(content: string, filePath: string): number {
-    const extension = filePath.split('.').pop()?.toLowerCase()
-    
-    // Adjust counting based on file type
-    switch (extension) {
-        case 'json':
-            // JSON is token-heavy due to strings and structure
-            return countTokens(content) * 1.1
-        case 'md':
-            // Markdown has more natural language
-            return countTokens(content) * 0.9
-        case 'ts':
-        case 'tsx':
-        case 'js':
-        case 'jsx':
-            // Code files benefit from pattern recognition
-            return countTokens(content)
-        default:
-            // Use standard counting for unknown types
-            return countTokens(content)
-    }
+    const ext = filePath.split('.').pop()?.toLowerCase()
+    if (ext === 'md') return Math.ceil(countTokens(content) * 0.9)
+    return countTokens(content)
 }
 
 /**
- * Token budget manager with overflow protection
+ * Token budget manager — tracks usage and truncates content to fit.
  */
 export class TokenBudget {
-    constructor(private maxTokens: number, private overflowAllowance: number = 0.1) {}
-    
-    /**
-     * Check if content fits within budget
-     */
+    private used = 0
+
+    constructor(
+        private readonly maxTokens: number,
+        private readonly overflowAllowance: number = 0.1,
+    ) {}
+
+    get remaining(): number {
+        return Math.max(0, this.maxTokens - this.used)
+    }
+
     fits(content: string): boolean {
-        const tokens = countTokens(content)
-        return tokens <= this.maxTokens * (1 + this.overflowAllowance)
+        return countTokensFast(content) <= this.remaining * (1 + this.overflowAllowance)
     }
-    
-    /**
-     * Get remaining token count
-     */
-    remaining(usedTokens: number): number {
-        return Math.max(0, this.maxTokens - usedTokens)
+
+    consume(tokens: number): boolean {
+        this.used += tokens
+        return this.used <= this.maxTokens * (1 + this.overflowAllowance)
     }
-    
-    /**
-     * Truncate content to fit within budget
-     */
-    truncate(content: string, usedTokens: number = 0): string {
-        const available = this.remaining(usedTokens)
-        if (available <= 0) return ''
-        
-        const estimatedTokens = countTokens(content)
-        if (estimatedTokens <= available) return content
-        
-        // Rough truncation based on character ratio
-        const ratio = available / estimatedTokens
-        const truncateAt = Math.floor(content.length * ratio * 0.9) // 10% buffer
-        
-        return content.substring(0, truncateAt) + '\n... [truncated due to token budget]'
+
+    truncate(content: string): string {
+        if (this.remaining <= 0) return ''
+        const estimated = countTokensFast(content)
+        if (estimated <= this.remaining) return content
+        const ratio = this.remaining / estimated
+        const cutAt = Math.floor(content.length * ratio * 0.9)
+        return content.slice(0, cutAt) + '\n… [truncated — token budget reached]'
     }
 }