diff --git a/src/core/graph.ts b/src/core/graph.ts index b66494e..6e82964 100644 --- a/src/core/graph.ts +++ b/src/core/graph.ts @@ -57,6 +57,12 @@ export function getConnections(nodeId: string): { incoming: (Edge & { node: Node }; } +export function getEdgesByNode(nodeId: string): Edge[] { + const db = getDb(); + const rows = db.prepare('SELECT * FROM edges WHERE from_id = ? OR to_id = ?').all(nodeId, nodeId) as any[]; + return rows.map(rowToEdge); +} + export function buildTree(rootId?: string): TreeNode[] { const db = getDb(); const edges = db.prepare('SELECT * FROM edges').all() as any[]; diff --git a/src/mcp/index.ts b/src/mcp/index.ts index 099ceac..09a20b6 100644 --- a/src/mcp/index.ts +++ b/src/mcp/index.ts @@ -1,8 +1,10 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { z } from 'zod/v3'; -import { query, listNodes, getNode, findNodeByPrefix, addNode, addEdge } from '../core/store'; -import { getConnections } from '../core/graph'; +import { query, listNodes, getNode, findNodeByPrefix, addNode, addEdge, removeNode, removeEdge, updateNode } from '../core/store'; +import { getConnections, getEdgesByNode } from '../core/graph'; +import { cosineSimilarity } from '../core/search/vector'; +import { getDb } from '../core/db'; import { NodeKind, EdgeType } from '../types'; function serialize(data: any): string { @@ -113,6 +115,256 @@ server.tool( } ); +// --- memory_split --- +server.tool( + 'memory_split', + 'Break a large node into smaller children, updating parent with a summary', + { + id: z.string().describe('Node ID to split'), + pieces: z.array(z.object({ title: z.string(), content: z.string() })).describe('Child pieces to create'), + summary: z.string().optional().describe('Optional summary to replace parent content'), + }, + async ({ id, pieces, summary }) => { + const node = getNode(id) ?? findNodeByPrefix(id); + if (!node) return { content: [{ type: 'text' as const, text: serialize({ error: 'Node not found' }) }], isError: true }; + + if (summary !== undefined) { + await updateNode(node.id, { content: summary }); + } + + const children = []; + for (const piece of pieces) { + const child = await addNode({ kind: node.kind, title: piece.title, content: piece.content, tags: [...node.tags] }); + addEdge(node.id, child.id, 'contains'); + children.push(child); + } + return { content: [{ type: 'text' as const, text: serialize({ parent: node.id, children: children.map(c => ({ id: c.id, title: c.title })) }) }] }; + } +); + +// --- memory_merge --- +server.tool( + 'memory_merge', + 'Merge multiple nodes into a single new node, relinking edges and deleting originals', + { + nodeIds: z.array(z.string()).describe('Node IDs to merge'), + title: z.string().describe('Title for merged node'), + content: z.string().describe('Content for merged node'), + kind: z.enum(['memory', 'component', 'task', 'decision']).optional().describe('Kind (defaults to first node\'s kind)'), + }, + async ({ nodeIds, title, content, kind }) => { + const nodes = nodeIds.map(id => getNode(id) ?? findNodeByPrefix(id)); + if (nodes.some(n => !n)) return { content: [{ type: 'text' as const, text: serialize({ error: 'One or more nodes not found' }) }], isError: true }; + + const validNodes = nodes as Exclude[]; + const allTags = [...new Set(validNodes.flatMap(n => n.tags))]; + const mergedKind = (kind as NodeKind) ?? validNodes[0].kind; + + const merged = await addNode({ kind: mergedKind, title, content, tags: allTags }); + const oldIds = new Set(validNodes.map(n => n.id)); + + // Relink edges from old nodes to merged node + const db = getDb(); + for (const n of validNodes) { + const edges = getEdgesByNode(n.id); + for (const e of edges) { + const from = oldIds.has(e.fromId) ? merged.id : e.fromId; + const to = oldIds.has(e.toId) ? merged.id : e.toId; + if (from !== to) { + try { addEdge(from, to, e.type, e.metadata); } catch {} + } + } + removeNode(n.id, true); + } + + return { content: [{ type: 'text' as const, text: serialize({ merged: { id: merged.id, title: merged.title, tags: merged.tags } }) }] }; + } +); + +// --- memory_dedupe --- +server.tool( + 'memory_dedupe', + 'Find similar/duplicate nodes using embedding cosine similarity', + { + threshold: z.number().optional().describe('Similarity threshold (default 0.85)'), + kind: z.enum(['memory', 'component', 'task', 'decision']).optional().describe('Filter by kind'), + limit: z.number().optional().describe('Max groups to return (default 10)'), + }, + async ({ threshold, kind, limit }) => { + const thresh = threshold ?? 0.85; + const maxGroups = limit ?? 10; + const nodes = listNodes({ kind: kind as NodeKind, includeStale: false }); + const withEmb = nodes.filter(n => n.embedding); + + const groups: { nodes: { id: string; title: string }[]; similarity: number }[] = []; + const seen = new Set(); + + for (let i = 0; i < withEmb.length && groups.length < maxGroups; i++) { + if (seen.has(withEmb[i].id)) continue; + const group = [{ id: withEmb[i].id, title: withEmb[i].title }]; + let maxSim = 0; + for (let j = i + 1; j < withEmb.length; j++) { + if (seen.has(withEmb[j].id)) continue; + const sim = cosineSimilarity(withEmb[i].embedding!, withEmb[j].embedding!); + if (sim >= thresh) { + group.push({ id: withEmb[j].id, title: withEmb[j].title }); + seen.add(withEmb[j].id); + maxSim = Math.max(maxSim, sim); + } + } + if (group.length > 1) { + seen.add(withEmb[i].id); + groups.push({ nodes: group, similarity: maxSim }); + } + } + + return { content: [{ type: 'text' as const, text: serialize({ groups, count: groups.length }) }] }; + } +); + +// --- memory_prune --- +server.tool( + 'memory_prune', + 'Clean up the graph: delete stale nodes, decay old nodes, or remove orphans', + { + mode: z.enum(['hard_delete_stale', 'decay_then_delete', 'delete_orphans']).describe('Prune mode'), + maxAgeDays: z.number().optional().describe('Max age in days for decay/delete modes'), + }, + async ({ mode, maxAgeDays }) => { + const db = getDb(); + let affected = 0; + + if (mode === 'hard_delete_stale') { + const result = db.prepare('DELETE FROM nodes WHERE is_stale = 1').run(); + affected = result.changes; + } else if (mode === 'decay_then_delete') { + const cutoff = Date.now() - (maxAgeDays ?? 90) * 86400000; + // First mark old non-stale as stale + const decayed = db.prepare('UPDATE nodes SET is_stale = 1, updated_at = ? WHERE is_stale = 0 AND last_accessed_at < ?').run(Date.now(), cutoff); + // Then hard-delete already-stale that are also old + const deleted = db.prepare('DELETE FROM nodes WHERE is_stale = 1 AND updated_at < ?').run(cutoff); + affected = decayed.changes + deleted.changes; + } else if (mode === 'delete_orphans') { + const orphans = db.prepare(` + SELECT n.id FROM nodes n + WHERE n.is_stale = 0 + AND NOT EXISTS (SELECT 1 FROM edges e WHERE e.from_id = n.id OR e.to_id = n.id) + `).all() as any[]; + for (const o of orphans) { + removeNode(o.id, true); + } + affected = orphans.length; + } + + return { content: [{ type: 'text' as const, text: serialize({ mode, affected }) }] }; + } +); + +// --- memory_reorganize --- +server.tool( + 'memory_reorganize', + 'Move a node under a new parent (updates contains edges)', + { + nodeId: z.string().describe('Node ID to move'), + newParentId: z.string().describe('New parent node ID'), + }, + async ({ nodeId, newParentId }) => { + const node = getNode(nodeId) ?? findNodeByPrefix(nodeId); + const parent = getNode(newParentId) ?? findNodeByPrefix(newParentId); + if (!node) return { content: [{ type: 'text' as const, text: serialize({ error: 'Node not found' }) }], isError: true }; + if (!parent) return { content: [{ type: 'text' as const, text: serialize({ error: 'Parent not found' }) }], isError: true }; + + // Remove existing incoming contains edges + const db = getDb(); + const incomingContains = db.prepare('SELECT id FROM edges WHERE to_id = ? AND type = ?').all(node.id, 'contains') as any[]; + for (const e of incomingContains) { + removeEdge(e.id); + } + + const edge = addEdge(parent.id, node.id, 'contains'); + return { content: [{ type: 'text' as const, text: serialize({ moved: node.id, newParent: parent.id, edge: edge.id }) }] }; + } +); + +// --- memory_bulk_tag --- +server.tool( + 'memory_bulk_tag', + 'Add or remove tags on multiple nodes at once', + { + action: z.enum(['add', 'remove']).describe('Whether to add or remove tags'), + tags: z.array(z.string()).describe('Tags to add/remove'), + nodeIds: z.array(z.string()).optional().describe('Specific node IDs'), + filter: z.object({ + kind: z.enum(['memory', 'component', 'task', 'decision']).optional(), + status: z.string().optional(), + tags: z.array(z.string()).optional(), + }).optional().describe('Filter to select nodes'), + }, + async ({ action, tags, nodeIds, filter }) => { + let targets: { id: string; tags: string[] }[]; + + if (nodeIds?.length) { + targets = nodeIds.map(id => { + const n = getNode(id) ?? findNodeByPrefix(id); + return n ? { id: n.id, tags: n.tags } : null; + }).filter(Boolean) as { id: string; tags: string[] }[]; + } else if (filter) { + targets = listNodes({ kind: filter.kind as NodeKind, status: filter.status, tags: filter.tags }).map(n => ({ id: n.id, tags: n.tags })); + } else { + return { content: [{ type: 'text' as const, text: serialize({ error: 'Provide nodeIds or filter' }) }], isError: true }; + } + + let modified = 0; + for (const t of targets) { + let newTags: string[]; + if (action === 'add') { + newTags = [...new Set([...t.tags, ...tags])]; + } else { + newTags = t.tags.filter(tag => !tags.includes(tag)); + } + if (JSON.stringify(newTags) !== JSON.stringify(t.tags)) { + await updateNode(t.id, { tags: newTags }); + modified++; + } + } + + return { content: [{ type: 'text' as const, text: serialize({ action, tags, modified, total: targets.length }) }] }; + } +); + +// --- memory_stats --- +server.tool( + 'memory_stats', + 'Get graph statistics: counts by kind, stale/orphan counts, tag distribution', + {}, + async () => { + const db = getDb(); + + const byKind = db.prepare('SELECT kind, COUNT(*) as count FROM nodes WHERE is_stale = 0 GROUP BY kind').all(); + const totalNodes = db.prepare('SELECT COUNT(*) as count FROM nodes WHERE is_stale = 0').get() as any; + const staleCount = db.prepare('SELECT COUNT(*) as count FROM nodes WHERE is_stale = 1').get() as any; + const edgeCount = db.prepare('SELECT COUNT(*) as count FROM edges').get() as any; + const edgesByType = db.prepare('SELECT type, COUNT(*) as count FROM edges GROUP BY type').all(); + const orphanCount = db.prepare(` + SELECT COUNT(*) as count FROM nodes n + WHERE n.is_stale = 0 + AND NOT EXISTS (SELECT 1 FROM edges e WHERE e.from_id = n.id OR e.to_id = n.id) + `).get() as any; + const tagDist = db.prepare('SELECT tag, COUNT(*) as count FROM node_tags GROUP BY tag ORDER BY count DESC LIMIT 20').all(); + + return { + content: [{ + type: 'text' as const, + text: serialize({ + nodes: { total: totalNodes.count, stale: staleCount.count, orphans: orphanCount.count, byKind }, + edges: { total: edgeCount.count, byType: edgesByType }, + topTags: tagDist, + }), + }], + }; + } +); + async function main() { const transport = new StdioServerTransport(); await server.connect(transport);