Files
claude-task-master/mcp-server/src/context-manager.js

874 lines
25 KiB
JavaScript

import { logger } from "../../scripts/modules/utils.js";
import fs from "fs/promises";
import path from "path";
import { fileURLToPath } from "url";
import crypto from "crypto";
import Fuse from "fuse.js";
// Constants
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
const CONTEXT_DIR =
process.env.MCP_CONTEXT_DIR || path.join(__dirname, "../contexts");
const MAX_CONTEXT_HISTORY = parseInt(
process.env.MCP_MAX_CONTEXT_HISTORY || "10",
10
);
/**
* Context Manager for MCP server
* Handles storage, retrieval, and manipulation of context data
* Implements efficient indexing, versioning, and advanced context operations
*/
class ContextManager {
constructor() {
this.contexts = new Map();
this.contextHistory = new Map(); // For version history
this.contextIndex = null; // For fuzzy search
this.logger = logger;
this.ensureContextDir();
this.rebuildSearchIndex();
}
/**
* Ensure the contexts directory exists
*/
async ensureContextDir() {
try {
await fs.mkdir(CONTEXT_DIR, { recursive: true });
this.logger.info(`Context directory ensured at ${CONTEXT_DIR}`);
// Also create a versions subdirectory for history
await fs.mkdir(path.join(CONTEXT_DIR, "versions"), { recursive: true });
} catch (error) {
this.logger.error(`Failed to create context directory: ${error.message}`);
throw error;
}
}
/**
* Rebuild the search index for efficient context lookup
*/
async rebuildSearchIndex() {
await this.loadAllContextsFromDisk();
const contextsForIndex = Array.from(this.contexts.values()).map((ctx) => ({
id: ctx.id,
content:
typeof ctx.data === "string" ? ctx.data : JSON.stringify(ctx.data),
tags: ctx.tags.join(" "),
metadata: Object.entries(ctx.metadata)
.map(([k, v]) => `${k}:${v}`)
.join(" "),
}));
this.contextIndex = new Fuse(contextsForIndex, {
keys: ["id", "content", "tags", "metadata"],
includeScore: true,
threshold: 0.6,
});
this.logger.info(
`Rebuilt search index with ${contextsForIndex.length} contexts`
);
}
/**
* Create a new context
* @param {string} contextId - Unique identifier for the context
* @param {object|string} contextData - Initial context data
* @param {object} metadata - Optional metadata for the context
* @returns {object} The created context
*/
async createContext(contextId, contextData, metadata = {}) {
if (this.contexts.has(contextId)) {
throw new Error(`Context with ID ${contextId} already exists`);
}
const timestamp = new Date().toISOString();
const versionId = this.generateVersionId();
const context = {
id: contextId,
data: contextData,
metadata: {
created: timestamp,
updated: timestamp,
version: versionId,
...metadata,
},
tags: metadata.tags || [],
size: this.estimateSize(contextData),
};
this.contexts.set(contextId, context);
// Initialize version history
this.contextHistory.set(contextId, [
{
versionId,
timestamp,
data: JSON.parse(JSON.stringify(contextData)), // Deep clone
metadata: { ...context.metadata },
},
]);
await this.persistContext(contextId);
await this.persistContextVersion(contextId, versionId);
// Update the search index
this.rebuildSearchIndex();
this.logger.info(`Created context: ${contextId} (version: ${versionId})`);
return context;
}
/**
* Retrieve a context by ID
* @param {string} contextId - The context ID to retrieve
* @param {string} versionId - Optional specific version to retrieve
* @returns {object} The context object
*/
async getContext(contextId, versionId = null) {
// If specific version requested, try to get it from history
if (versionId) {
return this.getContextVersion(contextId, versionId);
}
// Try to get from memory first
if (this.contexts.has(contextId)) {
return this.contexts.get(contextId);
}
// Try to load from disk
try {
const context = await this.loadContextFromDisk(contextId);
if (context) {
this.contexts.set(contextId, context);
return context;
}
} catch (error) {
this.logger.error(
`Failed to load context ${contextId}: ${error.message}`
);
}
throw new Error(`Context with ID ${contextId} not found`);
}
/**
* Get a specific version of a context
* @param {string} contextId - The context ID
* @param {string} versionId - The version ID
* @returns {object} The versioned context
*/
async getContextVersion(contextId, versionId) {
// Check if version history is in memory
if (this.contextHistory.has(contextId)) {
const history = this.contextHistory.get(contextId);
const version = history.find((v) => v.versionId === versionId);
if (version) {
return {
id: contextId,
data: version.data,
metadata: version.metadata,
tags: version.metadata.tags || [],
size: this.estimateSize(version.data),
versionId: version.versionId,
};
}
}
// Try to load from disk
try {
const versionPath = path.join(
CONTEXT_DIR,
"versions",
`${contextId}_${versionId}.json`
);
const data = await fs.readFile(versionPath, "utf8");
const version = JSON.parse(data);
// Add to memory cache
if (!this.contextHistory.has(contextId)) {
this.contextHistory.set(contextId, []);
}
const history = this.contextHistory.get(contextId);
history.push(version);
return {
id: contextId,
data: version.data,
metadata: version.metadata,
tags: version.metadata.tags || [],
size: this.estimateSize(version.data),
versionId: version.versionId,
};
} catch (error) {
this.logger.error(
`Failed to load context version ${contextId}@${versionId}: ${error.message}`
);
throw new Error(
`Context version ${versionId} for ${contextId} not found`
);
}
}
/**
* Update an existing context
* @param {string} contextId - The context ID to update
* @param {object|string} contextData - New context data
* @param {object} metadata - Optional metadata updates
* @param {boolean} createNewVersion - Whether to create a new version
* @returns {object} The updated context
*/
async updateContext(
contextId,
contextData,
metadata = {},
createNewVersion = true
) {
const context = await this.getContext(contextId);
const timestamp = new Date().toISOString();
// Generate a new version ID if requested
const versionId = createNewVersion
? this.generateVersionId()
: context.metadata.version;
// Create a backup of the current state for versioning
if (createNewVersion) {
// Store the current version in history
if (!this.contextHistory.has(contextId)) {
this.contextHistory.set(contextId, []);
}
const history = this.contextHistory.get(contextId);
// Add current state to history
history.push({
versionId: context.metadata.version,
timestamp: context.metadata.updated,
data: JSON.parse(JSON.stringify(context.data)), // Deep clone
metadata: { ...context.metadata },
});
// Trim history if it exceeds the maximum size
if (history.length > MAX_CONTEXT_HISTORY) {
const excessVersions = history.splice(
0,
history.length - MAX_CONTEXT_HISTORY
);
// Clean up excess versions from disk
for (const version of excessVersions) {
this.removeContextVersionFile(contextId, version.versionId).catch(
(err) =>
this.logger.error(
`Failed to remove old version file: ${err.message}`
)
);
}
}
// Persist version
await this.persistContextVersion(contextId, context.metadata.version);
}
// Update the context
context.data = contextData;
context.metadata = {
...context.metadata,
...metadata,
updated: timestamp,
};
if (createNewVersion) {
context.metadata.version = versionId;
context.metadata.previousVersion = context.metadata.version;
}
if (metadata.tags) {
context.tags = metadata.tags;
}
// Update size estimate
context.size = this.estimateSize(contextData);
this.contexts.set(contextId, context);
await this.persistContext(contextId);
// Update the search index
this.rebuildSearchIndex();
this.logger.info(`Updated context: ${contextId} (version: ${versionId})`);
return context;
}
/**
* Delete a context and all its versions
* @param {string} contextId - The context ID to delete
* @returns {boolean} True if deletion was successful
*/
async deleteContext(contextId) {
if (!this.contexts.has(contextId)) {
const contextPath = path.join(CONTEXT_DIR, `${contextId}.json`);
try {
await fs.access(contextPath);
} catch (error) {
throw new Error(`Context with ID ${contextId} not found`);
}
}
this.contexts.delete(contextId);
// Remove from history
const history = this.contextHistory.get(contextId) || [];
this.contextHistory.delete(contextId);
try {
// Delete main context file
const contextPath = path.join(CONTEXT_DIR, `${contextId}.json`);
await fs.unlink(contextPath);
// Delete all version files
for (const version of history) {
await this.removeContextVersionFile(contextId, version.versionId);
}
// Update the search index
this.rebuildSearchIndex();
this.logger.info(`Deleted context: ${contextId}`);
return true;
} catch (error) {
this.logger.error(
`Failed to delete context files for ${contextId}: ${error.message}`
);
throw error;
}
}
/**
* List all available contexts with pagination and advanced filtering
* @param {object} options - Options for listing contexts
* @param {object} options.filters - Filters to apply
* @param {number} options.limit - Maximum number of contexts to return
* @param {number} options.offset - Number of contexts to skip
* @param {string} options.sortBy - Field to sort by
* @param {string} options.sortDirection - Sort direction ('asc' or 'desc')
* @param {string} options.query - Free text search query
* @returns {Array} Array of context objects
*/
async listContexts(options = {}) {
// Load all contexts from disk first
await this.loadAllContextsFromDisk();
const {
filters = {},
limit = 100,
offset = 0,
sortBy = "updated",
sortDirection = "desc",
query = "",
} = options;
let contexts;
// If there's a search query, use the search index
if (query && this.contextIndex) {
const searchResults = this.contextIndex.search(query);
contexts = searchResults.map((result) =>
this.contexts.get(result.item.id)
);
} else {
contexts = Array.from(this.contexts.values());
}
// Apply filters
if (filters.tag) {
contexts = contexts.filter(
(ctx) => ctx.tags && ctx.tags.includes(filters.tag)
);
}
if (filters.metadataKey && filters.metadataValue) {
contexts = contexts.filter(
(ctx) =>
ctx.metadata &&
ctx.metadata[filters.metadataKey] === filters.metadataValue
);
}
if (filters.createdAfter) {
const timestamp = new Date(filters.createdAfter);
contexts = contexts.filter(
(ctx) => new Date(ctx.metadata.created) >= timestamp
);
}
if (filters.updatedAfter) {
const timestamp = new Date(filters.updatedAfter);
contexts = contexts.filter(
(ctx) => new Date(ctx.metadata.updated) >= timestamp
);
}
// Apply sorting
contexts.sort((a, b) => {
let valueA, valueB;
if (sortBy === "created" || sortBy === "updated") {
valueA = new Date(a.metadata[sortBy]).getTime();
valueB = new Date(b.metadata[sortBy]).getTime();
} else if (sortBy === "size") {
valueA = a.size || 0;
valueB = b.size || 0;
} else if (sortBy === "id") {
valueA = a.id;
valueB = b.id;
} else {
valueA = a.metadata[sortBy];
valueB = b.metadata[sortBy];
}
if (valueA === valueB) return 0;
const sortFactor = sortDirection === "asc" ? 1 : -1;
return valueA < valueB ? -1 * sortFactor : 1 * sortFactor;
});
// Apply pagination
const paginatedContexts = contexts.slice(offset, offset + limit);
return {
contexts: paginatedContexts,
total: contexts.length,
offset,
limit,
hasMore: offset + limit < contexts.length,
};
}
/**
* Get the version history of a context
* @param {string} contextId - The context ID
* @returns {Array} Array of version objects
*/
async getContextHistory(contextId) {
// Ensure context exists
await this.getContext(contextId);
// Load history if not in memory
if (!this.contextHistory.has(contextId)) {
await this.loadContextHistoryFromDisk(contextId);
}
const history = this.contextHistory.get(contextId) || [];
// Return versions in reverse chronological order (newest first)
return history.sort((a, b) => {
const timeA = new Date(a.timestamp).getTime();
const timeB = new Date(b.timestamp).getTime();
return timeB - timeA;
});
}
/**
* Add tags to a context
* @param {string} contextId - The context ID
* @param {Array} tags - Array of tags to add
* @returns {object} The updated context
*/
async addTags(contextId, tags) {
const context = await this.getContext(contextId);
const currentTags = context.tags || [];
const uniqueTags = [...new Set([...currentTags, ...tags])];
// Update context with new tags
return this.updateContext(
contextId,
context.data,
{
tags: uniqueTags,
},
false
); // Don't create a new version for tag updates
}
/**
* Remove tags from a context
* @param {string} contextId - The context ID
* @param {Array} tags - Array of tags to remove
* @returns {object} The updated context
*/
async removeTags(contextId, tags) {
const context = await this.getContext(contextId);
const currentTags = context.tags || [];
const newTags = currentTags.filter((tag) => !tags.includes(tag));
// Update context with new tags
return this.updateContext(
contextId,
context.data,
{
tags: newTags,
},
false
); // Don't create a new version for tag updates
}
/**
* Handle context windowing and truncation
* @param {string} contextId - The context ID
* @param {number} maxSize - Maximum size in tokens/chars
* @param {string} strategy - Truncation strategy ('start', 'end', 'middle')
* @returns {object} The truncated context
*/
async truncateContext(contextId, maxSize, strategy = "end") {
const context = await this.getContext(contextId);
const contextText =
typeof context.data === "string"
? context.data
: JSON.stringify(context.data);
if (contextText.length <= maxSize) {
return context; // No truncation needed
}
let truncatedData;
switch (strategy) {
case "start":
truncatedData = contextText.slice(contextText.length - maxSize);
break;
case "middle":
const halfSize = Math.floor(maxSize / 2);
truncatedData =
contextText.slice(0, halfSize) +
"...[truncated]..." +
contextText.slice(contextText.length - halfSize);
break;
case "end":
default:
truncatedData = contextText.slice(0, maxSize);
break;
}
// If original data was an object, try to parse the truncated data
// Otherwise use it as a string
let updatedData;
if (typeof context.data === "object") {
try {
// This may fail if truncation broke JSON structure
updatedData = {
...context.data,
truncated: true,
truncation_strategy: strategy,
original_size: contextText.length,
truncated_size: truncatedData.length,
};
} catch (error) {
updatedData = truncatedData;
}
} else {
updatedData = truncatedData;
}
// Update with truncated data
return this.updateContext(
contextId,
updatedData,
{
truncated: true,
truncation_strategy: strategy,
original_size: contextText.length,
truncated_size: truncatedData.length,
},
true
); // Create a new version for the truncated data
}
/**
* Merge multiple contexts into a new context
* @param {Array} contextIds - Array of context IDs to merge
* @param {string} newContextId - ID for the new merged context
* @param {object} metadata - Optional metadata for the new context
* @returns {object} The new merged context
*/
async mergeContexts(contextIds, newContextId, metadata = {}) {
if (contextIds.length === 0) {
throw new Error("At least one context ID must be provided for merging");
}
if (this.contexts.has(newContextId)) {
throw new Error(`Context with ID ${newContextId} already exists`);
}
// Load all contexts to be merged
const contextsToMerge = [];
for (const id of contextIds) {
try {
const context = await this.getContext(id);
contextsToMerge.push(context);
} catch (error) {
this.logger.error(
`Could not load context ${id} for merging: ${error.message}`
);
throw new Error(`Failed to merge contexts: ${error.message}`);
}
}
// Check data types and decide how to merge
const allStrings = contextsToMerge.every((c) => typeof c.data === "string");
const allObjects = contextsToMerge.every(
(c) => typeof c.data === "object" && c.data !== null
);
let mergedData;
if (allStrings) {
// Merge strings with newlines between them
mergedData = contextsToMerge.map((c) => c.data).join("\n\n");
} else if (allObjects) {
// Merge objects by combining their properties
mergedData = {};
for (const context of contextsToMerge) {
mergedData = { ...mergedData, ...context.data };
}
} else {
// Convert everything to strings and concatenate
mergedData = contextsToMerge
.map((c) =>
typeof c.data === "string" ? c.data : JSON.stringify(c.data)
)
.join("\n\n");
}
// Collect all tags from merged contexts
const allTags = new Set();
for (const context of contextsToMerge) {
for (const tag of context.tags || []) {
allTags.add(tag);
}
}
// Create merged metadata
const mergedMetadata = {
...metadata,
tags: [...allTags],
merged_from: contextIds,
merged_at: new Date().toISOString(),
};
// Create the new merged context
return this.createContext(newContextId, mergedData, mergedMetadata);
}
/**
* Persist a context to disk
* @param {string} contextId - The context ID to persist
* @returns {Promise<void>}
*/
async persistContext(contextId) {
const context = this.contexts.get(contextId);
if (!context) {
throw new Error(`Context with ID ${contextId} not found`);
}
const contextPath = path.join(CONTEXT_DIR, `${contextId}.json`);
try {
await fs.writeFile(contextPath, JSON.stringify(context, null, 2), "utf8");
this.logger.debug(`Persisted context ${contextId} to disk`);
} catch (error) {
this.logger.error(
`Failed to persist context ${contextId}: ${error.message}`
);
throw error;
}
}
/**
* Persist a context version to disk
* @param {string} contextId - The context ID
* @param {string} versionId - The version ID
* @returns {Promise<void>}
*/
async persistContextVersion(contextId, versionId) {
if (!this.contextHistory.has(contextId)) {
throw new Error(`Context history for ${contextId} not found`);
}
const history = this.contextHistory.get(contextId);
const version = history.find((v) => v.versionId === versionId);
if (!version) {
throw new Error(`Version ${versionId} of context ${contextId} not found`);
}
const versionPath = path.join(
CONTEXT_DIR,
"versions",
`${contextId}_${versionId}.json`
);
try {
await fs.writeFile(versionPath, JSON.stringify(version, null, 2), "utf8");
this.logger.debug(
`Persisted context version ${contextId}@${versionId} to disk`
);
} catch (error) {
this.logger.error(
`Failed to persist context version ${contextId}@${versionId}: ${error.message}`
);
throw error;
}
}
/**
* Remove a context version file from disk
* @param {string} contextId - The context ID
* @param {string} versionId - The version ID
* @returns {Promise<void>}
*/
async removeContextVersionFile(contextId, versionId) {
const versionPath = path.join(
CONTEXT_DIR,
"versions",
`${contextId}_${versionId}.json`
);
try {
await fs.unlink(versionPath);
this.logger.debug(
`Removed context version file ${contextId}@${versionId}`
);
} catch (error) {
if (error.code !== "ENOENT") {
this.logger.error(
`Failed to remove context version file ${contextId}@${versionId}: ${error.message}`
);
throw error;
}
}
}
/**
* Load a context from disk
* @param {string} contextId - The context ID to load
* @returns {Promise<object>} The loaded context
*/
async loadContextFromDisk(contextId) {
const contextPath = path.join(CONTEXT_DIR, `${contextId}.json`);
try {
const data = await fs.readFile(contextPath, "utf8");
const context = JSON.parse(data);
this.logger.debug(`Loaded context ${contextId} from disk`);
return context;
} catch (error) {
this.logger.error(
`Failed to load context ${contextId} from disk: ${error.message}`
);
throw error;
}
}
/**
* Load context history from disk
* @param {string} contextId - The context ID
* @returns {Promise<Array>} The loaded history
*/
async loadContextHistoryFromDisk(contextId) {
try {
const files = await fs.readdir(path.join(CONTEXT_DIR, "versions"));
const versionFiles = files.filter(
(file) => file.startsWith(`${contextId}_`) && file.endsWith(".json")
);
const history = [];
for (const file of versionFiles) {
try {
const data = await fs.readFile(
path.join(CONTEXT_DIR, "versions", file),
"utf8"
);
const version = JSON.parse(data);
history.push(version);
} catch (error) {
this.logger.error(
`Failed to load context version file ${file}: ${error.message}`
);
}
}
this.contextHistory.set(contextId, history);
this.logger.debug(
`Loaded ${history.length} versions for context ${contextId}`
);
return history;
} catch (error) {
this.logger.error(
`Failed to load context history for ${contextId}: ${error.message}`
);
this.contextHistory.set(contextId, []);
return [];
}
}
/**
* Load all contexts from disk
* @returns {Promise<void>}
*/
async loadAllContextsFromDisk() {
try {
const files = await fs.readdir(CONTEXT_DIR);
const contextFiles = files.filter((file) => file.endsWith(".json"));
for (const file of contextFiles) {
const contextId = path.basename(file, ".json");
if (!this.contexts.has(contextId)) {
try {
const context = await this.loadContextFromDisk(contextId);
this.contexts.set(contextId, context);
} catch (error) {
// Already logged in loadContextFromDisk
}
}
}
this.logger.info(`Loaded ${this.contexts.size} contexts from disk`);
} catch (error) {
this.logger.error(`Failed to load contexts from disk: ${error.message}`);
throw error;
}
}
/**
* Generate a unique version ID
* @returns {string} A unique version ID
*/
generateVersionId() {
return crypto.randomBytes(8).toString("hex");
}
/**
* Estimate the size of context data
* @param {object|string} data - The context data
* @returns {number} Estimated size in bytes
*/
estimateSize(data) {
if (typeof data === "string") {
return Buffer.byteLength(data, "utf8");
}
if (typeof data === "object" && data !== null) {
return Buffer.byteLength(JSON.stringify(data), "utf8");
}
return 0;
}
}
export default ContextManager;