@stackmemoryai/stackmemory
Version:
Project-scoped memory for AI coding tools. Durable context across sessions with MCP integration, frames, smart retrieval, Claude Code skills, and automatic hooks.
302 lines (301 loc) • 9.42 kB
JavaScript
import { fileURLToPath as __fileURLToPath } from 'url';
import { dirname as __pathDirname } from 'path';
const __filename = __fileURLToPath(import.meta.url);
const __dirname = __pathDirname(__filename);
import { logger } from "../../../core/monitoring/logger.js";
class ContextBudgetManager {
config;
tokenUsage = /* @__PURE__ */ new Map();
DEFAULT_MAX_TOKENS = 4e3;
TOKEN_CHAR_RATIO = 0.25;
// Rough estimate: 1 token ≈ 4 chars
constructor(config) {
this.config = {
maxTokens: config?.maxTokens || this.DEFAULT_MAX_TOKENS,
priorityWeights: {
task: config?.priorityWeights?.task || 0.3,
recentWork: config?.priorityWeights?.recentWork || 0.25,
feedback: config?.priorityWeights?.feedback || 0.2,
gitHistory: config?.priorityWeights?.gitHistory || 0.15,
dependencies: config?.priorityWeights?.dependencies || 0.1
},
compressionEnabled: config?.compressionEnabled ?? true,
adaptiveBudgeting: config?.adaptiveBudgeting ?? true
};
}
/**
* Estimate tokens for a given text
*/
estimateTokens(text) {
if (!text) return 0;
const baseTokens = text.length * this.TOKEN_CHAR_RATIO;
const codeMultiplier = this.detectCodeContent(text) ? 1.2 : 1;
const jsonMultiplier = this.detectJsonContent(text) ? 0.9 : 1;
return Math.ceil(baseTokens * codeMultiplier * jsonMultiplier);
}
/**
* Allocate token budget across different context categories
*/
allocateBudget(context) {
const currentTokens = this.calculateCurrentTokens(context);
if (currentTokens <= this.config.maxTokens) {
logger.debug("Context within budget", {
used: currentTokens,
max: this.config.maxTokens
});
return context;
}
logger.info("Context exceeds budget, optimizing...", {
current: currentTokens,
max: this.config.maxTokens
});
if (this.config.adaptiveBudgeting) {
return this.adaptiveBudgetAllocation(context, currentTokens);
}
return this.priorityBasedAllocation(context, currentTokens);
}
/**
* Compress context to fit within budget
*/
compressContext(context) {
if (!this.config.compressionEnabled) {
return context;
}
const compressed = {
...context,
task: this.compressTaskContext(context.task),
history: this.compressHistoryContext(context.history),
environment: this.compressEnvironmentContext(context.environment),
memory: this.compressMemoryContext(context.memory),
tokenCount: 0
};
compressed.tokenCount = this.calculateCurrentTokens(compressed);
logger.debug("Context compressed", {
original: context.tokenCount,
compressed: compressed.tokenCount,
reduction: `${Math.round((1 - compressed.tokenCount / context.tokenCount) * 100)}%`
});
return compressed;
}
/**
* Get current token usage statistics
*/
getUsage() {
const categories = {};
let totalUsed = 0;
for (const [category, tokens] of this.tokenUsage) {
categories[category] = tokens;
totalUsed += tokens;
}
return {
used: totalUsed,
available: this.config.maxTokens - totalUsed,
categories
};
}
/**
* Calculate current token count for context
*/
calculateCurrentTokens(context) {
this.tokenUsage.clear();
const taskTokens = this.estimateTokens(JSON.stringify(context.task));
const historyTokens = this.estimateTokens(JSON.stringify(context.history));
const envTokens = this.estimateTokens(JSON.stringify(context.environment));
const memoryTokens = this.estimateTokens(JSON.stringify(context.memory));
this.tokenUsage.set("task", taskTokens);
this.tokenUsage.set("history", historyTokens);
this.tokenUsage.set("environment", envTokens);
this.tokenUsage.set("memory", memoryTokens);
return taskTokens + historyTokens + envTokens + memoryTokens;
}
/**
* Adaptive budget allocation based on iteration phase
*/
adaptiveBudgetAllocation(context, currentTokens) {
const reductionRatio = this.config.maxTokens / currentTokens;
const phase = this.determinePhase(context.task.currentIteration);
const adjustedWeights = this.getPhaseAdjustedWeights(phase);
return this.applyWeightedReduction(context, reductionRatio, adjustedWeights);
}
/**
* Priority-based allocation using fixed weights
*/
priorityBasedAllocation(context, currentTokens) {
const reductionRatio = this.config.maxTokens / currentTokens;
return this.applyWeightedReduction(context, reductionRatio, this.config.priorityWeights);
}
/**
* Apply weighted reduction to context
*/
applyWeightedReduction(context, reductionRatio, weights) {
const reduced = { ...context };
if (weights.recentWork < 1) {
const keepCount = Math.ceil(
context.history.recentIterations.length * reductionRatio * weights.recentWork
);
reduced.history = {
...context.history,
recentIterations: context.history.recentIterations.slice(-keepCount)
};
}
if (weights.gitHistory < 1) {
const keepCount = Math.ceil(
context.history.gitCommits.length * reductionRatio * weights.gitHistory
);
reduced.history.gitCommits = context.history.gitCommits.slice(-keepCount);
}
if (weights.dependencies < 1) {
const keepCount = Math.ceil(
context.memory.relevantFrames.length * reductionRatio * weights.dependencies
);
reduced.memory = {
...context.memory,
relevantFrames: context.memory.relevantFrames.sort((a, b) => (b.created_at || 0) - (a.created_at || 0)).slice(0, keepCount)
};
}
reduced.tokenCount = this.calculateCurrentTokens(reduced);
return reduced;
}
/**
* Compress task context
*/
compressTaskContext(task) {
return {
...task,
description: this.truncateWithEllipsis(task.description, 500),
criteria: task.criteria.slice(0, 5),
// Keep top 5 criteria
feedback: task.feedback ? this.truncateWithEllipsis(task.feedback, 300) : void 0
};
}
/**
* Compress history context
*/
compressHistoryContext(history) {
return {
...history,
recentIterations: history.recentIterations.slice(-5).map((iter) => ({
...iter,
summary: this.truncateWithEllipsis(iter.summary, 100)
})),
gitCommits: history.gitCommits.slice(-10).map((commit) => ({
...commit,
message: this.truncateWithEllipsis(commit.message, 80),
files: commit.files.slice(0, 5)
// Keep top 5 files
})),
changedFiles: history.changedFiles.slice(0, 20),
// Keep top 20 files
testResults: history.testResults.slice(-3)
// Keep last 3 test runs
};
}
/**
* Compress environment context
*/
compressEnvironmentContext(env) {
return {
...env,
dependencies: this.compressObject(env.dependencies, 20),
// Keep top 20 deps
configuration: this.compressObject(env.configuration, 10)
// Keep top 10 config items
};
}
/**
* Compress memory context
*/
compressMemoryContext(memory) {
return {
...memory,
relevantFrames: memory.relevantFrames.slice(0, 5),
// Keep top 5 frames
decisions: memory.decisions.filter((d) => d.impact !== "low").slice(-5),
// Keep last 5
patterns: memory.patterns.filter((p) => p.successRate > 0.7).slice(0, 3),
// Keep top 3
blockers: memory.blockers.filter((b) => !b.resolved)
// Keep unresolved only
};
}
/**
* Determine iteration phase
*/
determinePhase(iteration) {
if (iteration <= 3) return "early";
if (iteration <= 10) return "middle";
return "late";
}
/**
* Get phase-adjusted weights
*/
getPhaseAdjustedWeights(phase) {
switch (phase) {
case "early":
return {
task: 0.4,
recentWork: 0.1,
feedback: 0.2,
gitHistory: 0.2,
dependencies: 0.1
};
case "middle":
return this.config.priorityWeights;
case "late":
return {
task: 0.2,
recentWork: 0.35,
feedback: 0.25,
gitHistory: 0.15,
dependencies: 0.05
};
}
}
/**
* Detect if text contains code
*/
detectCodeContent(text) {
const codePatterns = [
/function\s+\w+\s*\(/,
/class\s+\w+/,
/const\s+\w+\s*=/,
/import\s+.*from/,
/\{[\s\S]*\}/
];
return codePatterns.some((pattern) => pattern.test(text));
}
/**
* Detect if text contains JSON
*/
detectJsonContent(text) {
try {
JSON.parse(text);
return true;
} catch {
return text.includes('"') && text.includes(":") && text.includes("{");
}
}
/**
* Truncate text with ellipsis
*/
truncateWithEllipsis(text, maxLength) {
if (text.length <= maxLength) return text;
return text.substring(0, maxLength - 3) + "...";
}
/**
* Compress object by keeping only top N entries
*/
compressObject(obj, maxEntries) {
const entries = Object.entries(obj);
if (entries.length <= maxEntries) return obj;
const compressed = {};
entries.slice(0, maxEntries).forEach(([key, value]) => {
compressed[key] = value;
});
return compressed;
}
}
export {
ContextBudgetManager
};
//# sourceMappingURL=context-budget-manager.js.map