agent-working-memory 0.3.0

package/src/mcp.ts

@@ -0,0 +1,635 @@
+/**
+ * MCP Server — Model Context Protocol interface for AgentWorkingMemory.
+ *
+ * Runs as a stdio-based MCP server that Claude Code connects to directly.
+ * Uses the storage and engine layers in-process (no HTTP overhead).
+ *
+ * Tools exposed (11):
+ *   memory_write       — store a memory (salience filter decides disposition)
+ *   memory_recall      — activate memories by context (cognitive retrieval)
+ *   memory_feedback    — report whether a recalled memory was useful
+ *   memory_retract     — invalidate a wrong memory with optional correction
+ *   memory_stats       — get memory health metrics
+ *   memory_checkpoint  — save structured execution state (survives compaction)
+ *   memory_restore     — restore state + targeted recall after compaction
+ *   memory_task_add    — create a prioritized task
+ *   memory_task_update — change task status, priority, or blocking
+ *   memory_task_list   — list tasks filtered by status
+ *   memory_task_next   — get the highest-priority actionable task
+ *
+ * Run: npx tsx src/mcp.ts
+ * Config: add to ~/.claude.json or .mcp.json
+ */
+
+import { readFileSync } from 'node:fs';
+import { resolve } from 'node:path';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+
+// Load .env file if present (no external dependency)
+try {
+  const envPath = resolve(process.cwd(), '.env');
+  const envContent = readFileSync(envPath, 'utf-8');
+  for (const line of envContent.split('\n')) {
+    const trimmed = line.trim();
+    if (!trimmed || trimmed.startsWith('#')) continue;
+    const eqIdx = trimmed.indexOf('=');
+    if (eqIdx === -1) continue;
+    const key = trimmed.slice(0, eqIdx).trim();
+    const val = trimmed.slice(eqIdx + 1).trim().replace(/^["']|["']$/g, '');
+    if (!process.env[key]) process.env[key] = val;
+  }
+} catch { /* No .env file */ }
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { z } from 'zod';
+
+import { EngramStore } from './storage/sqlite.js';
+import { ActivationEngine } from './engine/activation.js';
+import { ConnectionEngine } from './engine/connections.js';
+import { StagingBuffer } from './engine/staging.js';
+import { EvictionEngine } from './engine/eviction.js';
+import { RetractionEngine } from './engine/retraction.js';
+import { EvalEngine } from './engine/eval.js';
+import { ConsolidationEngine } from './engine/consolidation.js';
+import { ConsolidationScheduler } from './engine/consolidation-scheduler.js';
+import { evaluateSalience } from './core/salience.js';
+import type { ConsciousState } from './types/checkpoint.js';
+import type { SalienceEventType } from './core/salience.js';
+import type { TaskStatus, TaskPriority } from './types/engram.js';
+import { DEFAULT_AGENT_CONFIG } from './types/agent.js';
+import { embed } from './core/embeddings.js';
+
+// --- Setup ---
+
+const DB_PATH = process.env.AWM_DB_PATH ?? 'memory.db';
+const AGENT_ID = process.env.AWM_AGENT_ID ?? 'claude-code';
+
+const store = new EngramStore(DB_PATH);
+const activationEngine = new ActivationEngine(store);
+const connectionEngine = new ConnectionEngine(store, activationEngine);
+const stagingBuffer = new StagingBuffer(store, activationEngine);
+const evictionEngine = new EvictionEngine(store);
+const retractionEngine = new RetractionEngine(store);
+const evalEngine = new EvalEngine(store);
+const consolidationEngine = new ConsolidationEngine(store);
+const consolidationScheduler = new ConsolidationScheduler(store, consolidationEngine);
+
+stagingBuffer.start(DEFAULT_AGENT_CONFIG.stagingTtlMs);
+consolidationScheduler.start();
+
+const server = new McpServer({
+  name: 'agent-working-memory',
+  version: '0.3.0',
+});
+
+// --- Tools ---
+
+server.tool(
+  'memory_write',
+  `Store a memory. The salience filter decides whether it's worth keeping (active), needs more evidence (staging), or should be discarded.
+
+Use this when you learn something that might be useful later:
+- Discoveries about the codebase, bugs, or architecture
+- Decisions you made and why
+- Errors encountered and how they were resolved
+- User preferences or project patterns
+
+The concept should be a short label (3-8 words). The content should be the full detail.`,
+  {
+    concept: z.string().describe('Short label for this memory (3-8 words)'),
+    content: z.string().describe('Full detail of what was learned'),
+    tags: z.array(z.string()).optional().describe('Optional tags for categorization'),
+    event_type: z.enum(['observation', 'decision', 'friction', 'surprise', 'causal'])
+      .optional().default('observation')
+      .describe('Type of event: observation (default), decision, friction (error/blocker), surprise, causal (root cause)'),
+    surprise: z.number().min(0).max(1).optional().default(0.3)
+      .describe('How surprising was this? 0=expected, 1=very unexpected'),
+    decision_made: z.boolean().optional().default(false)
+      .describe('Was a decision made? True boosts importance'),
+    causal_depth: z.number().min(0).max(1).optional().default(0.3)
+      .describe('How deep is the causal understanding? 0=surface, 1=root cause'),
+    resolution_effort: z.number().min(0).max(1).optional().default(0.3)
+      .describe('How much effort to resolve? 0=trivial, 1=significant debugging'),
+  },
+  async (params) => {
+    const salience = evaluateSalience({
+      content: params.content,
+      eventType: params.event_type as SalienceEventType,
+      surprise: params.surprise,
+      decisionMade: params.decision_made,
+      causalDepth: params.causal_depth,
+      resolutionEffort: params.resolution_effort,
+    });
+
+    if (salience.disposition === 'discard') {
+      return {
+        content: [{
+          type: 'text' as const,
+          text: `Memory discarded (salience ${salience.score.toFixed(2)} below threshold). Not worth storing.`,
+        }],
+      };
+    }
+
+    const engram = store.createEngram({
+      agentId: AGENT_ID,
+      concept: params.concept,
+      content: params.content,
+      tags: params.tags,
+      salience: salience.score,
+      salienceFeatures: salience.features,
+      reasonCodes: salience.reasonCodes,
+      ttl: salience.disposition === 'staging' ? DEFAULT_AGENT_CONFIG.stagingTtlMs : undefined,
+    });
+
+    if (salience.disposition === 'staging') {
+      store.updateStage(engram.id, 'staging');
+    } else {
+      connectionEngine.enqueue(engram.id);
+    }
+
+    // Generate embedding asynchronously (don't block response)
+    embed(`${params.concept} ${params.content}`).then(vec => {
+      store.updateEmbedding(engram.id, vec);
+    }).catch(() => {}); // Embedding failure is non-fatal
+
+    // Auto-checkpoint: track write
+    try { store.updateAutoCheckpointWrite(AGENT_ID, engram.id); } catch { /* non-fatal */ }
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Memory stored (${salience.disposition}). ID: ${engram.id}\nConcept: ${params.concept}\nSalience: ${salience.score.toFixed(2)}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_recall',
+  `Recall memories relevant to a context. Uses cognitive activation — not keyword search.
+
+Use this when you need to remember something:
+- Starting a new task (recall relevant past experience)
+- Debugging (recall similar errors and solutions)
+- Making decisions (recall past decisions and outcomes)
+
+Returns the most relevant memories ranked by a composite score of text relevance, temporal recency, and associative strength.`,
+  {
+    context: z.string().describe('What are you thinking about? Describe the current situation or question'),
+    limit: z.number().optional().default(5).describe('Max memories to return (default 5)'),
+    min_score: z.number().optional().default(0.05).describe('Minimum relevance score (default 0.05)'),
+    include_staging: z.boolean().optional().default(false).describe('Include weak/unconfirmed memories?'),
+    use_reranker: z.boolean().optional().default(true).describe('Use cross-encoder re-ranking for better relevance (default true)'),
+    use_expansion: z.boolean().optional().default(true).describe('Expand query with synonyms for better recall (default true)'),
+  },
+  async (params) => {
+    const results = await activationEngine.activate({
+      agentId: AGENT_ID,
+      context: params.context,
+      limit: params.limit,
+      minScore: params.min_score,
+      includeStaging: params.include_staging,
+      useReranker: params.use_reranker,
+      useExpansion: params.use_expansion,
+    });
+
+    // Auto-checkpoint: track recall
+    try {
+      const ids = results.map(r => r.engram.id);
+      store.updateAutoCheckpointRecall(AGENT_ID, params.context, ids);
+    } catch { /* non-fatal */ }
+
+    if (results.length === 0) {
+      return {
+        content: [{
+          type: 'text' as const,
+          text: 'No relevant memories found.',
+        }],
+      };
+    }
+
+    const lines = results.map((r, i) => {
+      const tags = r.engram.tags?.length ? ` [${r.engram.tags.join(', ')}]` : '';
+      return `${i + 1}. **${r.engram.concept}** (score: ${r.score.toFixed(3)})${tags}\n   ${r.engram.content}\n   _${r.why}_\n   ID: ${r.engram.id}`;
+    });
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Recalled ${results.length} memories:\n\n${lines.join('\n\n')}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_feedback',
+  `Report whether a recalled memory was actually useful. This updates the memory's confidence score — useful memories become stronger, useless ones weaken.
+
+Always call this after using a recalled memory so the system learns what's valuable.`,
+  {
+    engram_id: z.string().describe('ID of the memory (from memory_recall results)'),
+    useful: z.boolean().describe('Was this memory actually helpful?'),
+    context: z.string().optional().describe('Brief note on why it was/wasn\'t useful'),
+  },
+  async (params) => {
+    store.logRetrievalFeedback(null, params.engram_id, params.useful, params.context ?? '');
+
+    const engram = store.getEngram(params.engram_id);
+    if (engram) {
+      const delta = params.useful
+        ? DEFAULT_AGENT_CONFIG.feedbackPositiveBoost
+        : -DEFAULT_AGENT_CONFIG.feedbackNegativePenalty;
+      store.updateConfidence(engram.id, engram.confidence + delta);
+    }
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Feedback recorded: ${params.useful ? 'useful' : 'not useful'}. Confidence ${params.useful ? 'increased' : 'decreased'}.`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_retract',
+  `Retract a memory that turned out to be wrong. Creates a correction and reduces confidence of related memories.
+
+Use this when you discover a memory contains incorrect information.`,
+  {
+    engram_id: z.string().describe('ID of the wrong memory'),
+    reason: z.string().describe('Why is this memory wrong?'),
+    correction: z.string().optional().describe('What is the correct information? (creates a new memory)'),
+  },
+  async (params) => {
+    const result = retractionEngine.retract({
+      agentId: AGENT_ID,
+      targetEngramId: params.engram_id,
+      reason: params.reason,
+      counterContent: params.correction,
+    });
+
+    const parts = [`Memory ${params.engram_id} retracted.`];
+    if (result.correctionId) {
+      parts.push(`Correction stored as ${result.correctionId}.`);
+    }
+    parts.push(`${result.associatesAffected} related memories had confidence reduced.`);
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: parts.join(' '),
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_stats',
+  `Get memory health stats — how many memories, confidence levels, association count, and system performance.`,
+  {},
+  async () => {
+    const metrics = evalEngine.computeMetrics(AGENT_ID);
+    const lines = [
+      `Agent: ${AGENT_ID}`,
+      `Active memories: ${metrics.activeEngramCount}`,
+      `Staging: ${metrics.stagingEngramCount}`,
+      `Retracted: ${metrics.retractedCount}`,
+      `Avg confidence: ${metrics.avgConfidence.toFixed(3)}`,
+      `Total edges: ${metrics.totalEdges}`,
+      `Edge utility: ${(metrics.edgeUtilityRate * 100).toFixed(1)}%`,
+      `Activations (24h): ${metrics.activationCount}`,
+      `Avg latency: ${metrics.avgLatencyMs.toFixed(1)}ms`,
+    ];
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: lines.join('\n'),
+      }],
+    };
+  }
+);
+
+// --- Checkpointing Tools ---
+
+server.tool(
+  'memory_checkpoint',
+  `Save your current execution state so you can recover after context compaction.
+
+Use this when:
+- You're about to do something that might fill the context window
+- You've made important decisions you don't want to lose
+- You want to preserve your working state before a long operation
+
+The state is saved per-agent and overwrites any previous checkpoint.`,
+  {
+    current_task: z.string().describe('What you are currently working on'),
+    decisions: z.array(z.string()).optional().default([])
+      .describe('Key decisions made so far'),
+    active_files: z.array(z.string()).optional().default([])
+      .describe('Files you are currently working with'),
+    next_steps: z.array(z.string()).optional().default([])
+      .describe('What needs to happen next'),
+    related_memory_ids: z.array(z.string()).optional().default([])
+      .describe('IDs of memories relevant to current work'),
+    notes: z.string().optional().default('')
+      .describe('Any other context worth preserving'),
+    episode_id: z.string().optional()
+      .describe('Current episode ID if known'),
+  },
+  async (params) => {
+    const state: ConsciousState = {
+      currentTask: params.current_task,
+      decisions: params.decisions,
+      activeFiles: params.active_files,
+      nextSteps: params.next_steps,
+      relatedMemoryIds: params.related_memory_ids,
+      notes: params.notes,
+      episodeId: params.episode_id ?? null,
+    };
+
+    store.saveCheckpoint(AGENT_ID, state);
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Checkpoint saved.\nTask: ${params.current_task}\nDecisions: ${params.decisions.length}\nNext steps: ${params.next_steps.length}\nFiles: ${params.active_files.length}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_restore',
+  `Restore your previous execution state after context compaction or at session start.
+
+Returns:
+- Your saved execution state (task, decisions, next steps, files)
+- Recently recalled memories for context
+- Your last write for continuity
+- How long you were idle
+
+Use this at the start of every session or after compaction to pick up where you left off.`,
+  {},
+  async () => {
+    const checkpoint = store.getCheckpoint(AGENT_ID);
+
+    const now = Date.now();
+    const idleMs = checkpoint
+      ? now - checkpoint.auto.lastActivityAt.getTime()
+      : 0;
+
+    // Get last written engram
+    let lastWrite: { id: string; concept: string; content: string } | null = null;
+    if (checkpoint?.auto.lastWriteId) {
+      const engram = store.getEngram(checkpoint.auto.lastWriteId);
+      if (engram) {
+        lastWrite = { id: engram.id, concept: engram.concept, content: engram.content };
+      }
+    }
+
+    // Recall memories using last context
+    let recalledMemories: Array<{ id: string; concept: string; content: string; score: number }> = [];
+    const recallContext = checkpoint?.auto.lastRecallContext
+      ?? checkpoint?.executionState?.currentTask
+      ?? null;
+
+    if (recallContext) {
+      try {
+        const results = await activationEngine.activate({
+          agentId: AGENT_ID,
+          context: recallContext,
+          limit: 5,
+          minScore: 0.05,
+          useReranker: true,
+          useExpansion: true,
+        });
+        recalledMemories = results.map(r => ({
+          id: r.engram.id,
+          concept: r.engram.concept,
+          content: r.engram.content,
+          score: r.score,
+        }));
+      } catch { /* recall failure is non-fatal */ }
+    }
+
+    // Trigger mini-consolidation if idle >5min
+    const MINI_IDLE_MS = 5 * 60_000;
+    let miniConsolidationTriggered = false;
+    if (idleMs > MINI_IDLE_MS) {
+      miniConsolidationTriggered = true;
+      consolidationScheduler.runMiniConsolidation(AGENT_ID).catch(() => {});
+    }
+
+    // Format response
+    const parts: string[] = [];
+    const idleMin = Math.round(idleMs / 60_000);
+    parts.push(`Idle: ${idleMin}min${miniConsolidationTriggered ? ' (mini-consolidation triggered)' : ''}`);
+
+    if (checkpoint?.executionState) {
+      const s = checkpoint.executionState;
+      parts.push(`\n**Current task:** ${s.currentTask}`);
+      if (s.decisions.length) parts.push(`**Decisions:** ${s.decisions.join('; ')}`);
+      if (s.nextSteps.length) parts.push(`**Next steps:** ${s.nextSteps.map((st, i) => `${i + 1}. ${st}`).join(', ')}`);
+      if (s.activeFiles.length) parts.push(`**Active files:** ${s.activeFiles.join(', ')}`);
+      if (s.notes) parts.push(`**Notes:** ${s.notes}`);
+      if (checkpoint.checkpointAt) parts.push(`_Saved at: ${checkpoint.checkpointAt.toISOString()}_`);
+    } else {
+      parts.push('\nNo explicit checkpoint saved.');
+    }
+
+    if (lastWrite) {
+      parts.push(`\n**Last write:** ${lastWrite.concept}\n${lastWrite.content}`);
+    }
+
+    if (recalledMemories.length > 0) {
+      parts.push(`\n**Recalled memories (${recalledMemories.length}):**`);
+      for (const m of recalledMemories) {
+        parts.push(`- **${m.concept}** (${m.score.toFixed(3)}): ${m.content.slice(0, 150)}${m.content.length > 150 ? '...' : ''}`);
+      }
+    }
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: parts.join('\n'),
+      }],
+    };
+  }
+);
+
+// --- Task Management Tools ---
+
+server.tool(
+  'memory_task_add',
+  `Create a task that you need to come back to. Tasks are memories with status and priority tracking.
+
+Use this when:
+- You identify work that needs doing but can't do it right now
+- The user mentions something to do later
+- You want to park a sub-task while focusing on something more urgent
+
+Tasks automatically get high salience so they won't be discarded.`,
+  {
+    concept: z.string().describe('Short task title (3-10 words)'),
+    content: z.string().describe('Full task description — what needs doing, context, acceptance criteria'),
+    tags: z.array(z.string()).optional().describe('Tags for categorization'),
+    priority: z.enum(['urgent', 'high', 'medium', 'low']).default('medium')
+      .describe('Task priority: urgent (do now), high (do soon), medium (normal), low (backlog)'),
+    blocked_by: z.string().optional().describe('ID of a task that must finish first'),
+  },
+  async (params) => {
+    const engram = store.createEngram({
+      agentId: AGENT_ID,
+      concept: params.concept,
+      content: params.content,
+      tags: [...(params.tags ?? []), 'task'],
+      salience: 0.9, // Tasks always high salience
+      confidence: 0.8,
+      salienceFeatures: {
+        surprise: 0.5,
+        decisionMade: true,
+        causalDepth: 0.5,
+        resolutionEffort: 0.5,
+        eventType: 'decision',
+      },
+      reasonCodes: ['task-created'],
+      taskStatus: params.blocked_by ? 'blocked' : 'open',
+      taskPriority: params.priority as TaskPriority,
+      blockedBy: params.blocked_by,
+    });
+
+    connectionEngine.enqueue(engram.id);
+
+    // Generate embedding asynchronously
+    embed(`${params.concept} ${params.content}`).then(vec => {
+      store.updateEmbedding(engram.id, vec);
+    }).catch(() => {});
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Task created: ${engram.id}\nTitle: ${params.concept}\nPriority: ${params.priority}\nStatus: ${engram.taskStatus}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_task_update',
+  `Update a task's status or priority. Use this to:
+- Start working on a task (open → in_progress)
+- Mark a task done (→ done)
+- Block a task on another (→ blocked)
+- Reprioritize (change priority)
+- Unblock a task (clear blocked_by)`,
+  {
+    task_id: z.string().describe('ID of the task to update'),
+    status: z.enum(['open', 'in_progress', 'blocked', 'done']).optional()
+      .describe('New status'),
+    priority: z.enum(['urgent', 'high', 'medium', 'low']).optional()
+      .describe('New priority'),
+    blocked_by: z.string().optional().describe('ID of blocking task (set to empty string to unblock)'),
+  },
+  async (params) => {
+    const engram = store.getEngram(params.task_id);
+    if (!engram || !engram.taskStatus) {
+      return { content: [{ type: 'text' as const, text: `Task not found: ${params.task_id}` }] };
+    }
+
+    if (params.blocked_by !== undefined) {
+      store.updateBlockedBy(params.task_id, params.blocked_by || null);
+    }
+    if (params.status) {
+      store.updateTaskStatus(params.task_id, params.status as TaskStatus);
+    }
+    if (params.priority) {
+      store.updateTaskPriority(params.task_id, params.priority as TaskPriority);
+    }
+
+    const updated = store.getEngram(params.task_id)!;
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Task updated: ${updated.concept}\nStatus: ${updated.taskStatus}\nPriority: ${updated.taskPriority}${updated.blockedBy ? `\nBlocked by: ${updated.blockedBy}` : ''}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_task_list',
+  `List tasks with optional status filter. Shows tasks ordered by priority (urgent first).
+
+Use at the start of a session to see what's pending, or to check blocked/done tasks.`,
+  {
+    status: z.enum(['open', 'in_progress', 'blocked', 'done']).optional()
+      .describe('Filter by status (omit to see all active tasks)'),
+    include_done: z.boolean().optional().default(false)
+      .describe('Include completed tasks?'),
+  },
+  async (params) => {
+    let tasks = store.getTasks(AGENT_ID, params.status as TaskStatus | undefined);
+    if (!params.include_done && !params.status) {
+      tasks = tasks.filter(t => t.taskStatus !== 'done');
+    }
+
+    if (tasks.length === 0) {
+      return { content: [{ type: 'text' as const, text: 'No tasks found.' }] };
+    }
+
+    const lines = tasks.map((t, i) => {
+      const blocked = t.blockedBy ? ` [blocked by ${t.blockedBy}]` : '';
+      const tags = t.tags?.filter(tag => tag !== 'task').join(', ');
+      return `${i + 1}. [${t.taskStatus}] **${t.concept}** (${t.taskPriority})${blocked}\n   ${t.content.slice(0, 120)}${t.content.length > 120 ? '...' : ''}\n   ${tags ? `Tags: ${tags} | ` : ''}ID: ${t.id}`;
+    });
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Tasks (${tasks.length}):\n\n${lines.join('\n\n')}`,
+      }],
+    };
+  }
+);
+
+server.tool(
+  'memory_task_next',
+  `Get the single most important task to work on next.
+
+Prioritizes: in_progress tasks first (finish what you started), then by priority level, then oldest first. Skips blocked and done tasks.
+
+Use this when you finish a task or need to decide what to do next.`,
+  {},
+  async () => {
+    const next = store.getNextTask(AGENT_ID);
+    if (!next) {
+      return { content: [{ type: 'text' as const, text: 'No actionable tasks. All clear!' }] };
+    }
+
+    const blocked = next.blockedBy ? `\nBlocked by: ${next.blockedBy}` : '';
+    const tags = next.tags?.filter(tag => tag !== 'task').join(', ');
+
+    return {
+      content: [{
+        type: 'text' as const,
+        text: `Next task:\n**${next.concept}** (${next.taskPriority})\nStatus: ${next.taskStatus}\n${next.content}${blocked}\n${tags ? `Tags: ${tags}\n` : ''}ID: ${next.id}`,
+      }],
+    };
+  }
+);
+
+// --- Start ---
+
+async function main() {
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+  // Log to stderr (stdout is reserved for MCP protocol)
+  console.error(`AgentWorkingMemory MCP server started (agent: ${AGENT_ID}, db: ${DB_PATH})`);
+}
+
+main().catch(err => {
+  console.error('MCP server failed:', err);
+  process.exit(1);
+});

package/src/storage/index.ts

@@ -0,0 +1 @@
+export * from './sqlite.js';