code-graph-context 2.8.0 → 2.10.0

package/dist/mcp/tools/swarm-message.tool.js

@@ -0,0 +1,362 @@
+/**
+ * Swarm Message Tool
+ * Direct agent-to-agent messaging via Neo4j for reliable coordination.
+ *
+ * Unlike pheromones (passive, decay-based stigmergy), messages are explicit
+ * and delivered to agents when they claim tasks. This ensures critical
+ * coordination signals (blocked, conflict, findings) are reliably received.
+ */
+import { z } from 'zod';
+import { Neo4jService } from '../../storage/neo4j/neo4j.service.js';
+import { TOOL_NAMES, TOOL_METADATA } from '../constants.js';
+import { createErrorResponse, createSuccessResponse, resolveProjectIdOrError, debugLog } from '../utils.js';
+import { MESSAGE_CATEGORY_KEYS, MESSAGE_DEFAULT_TTL_MS, generateMessageId } from './swarm-constants.js';
+// ============================================================================
+// NEO4J QUERIES
+// ============================================================================
+/**
+ * Send a message. Creates a SwarmMessage node with optional target agent.
+ * Broadcast messages (no toAgentId) are visible to all agents in the swarm.
+ */
+const SEND_MESSAGE_QUERY = `
+  CREATE (m:SwarmMessage {
+    id: $messageId,
+    projectId: $projectId,
+    swarmId: $swarmId,
+    fromAgentId: $fromAgentId,
+    toAgentId: $toAgentId,
+    category: $category,
+    content: $content,
+    taskId: $taskId,
+    filePaths: $filePaths,
+    timestamp: timestamp(),
+    expiresAt: timestamp() + $ttlMs,
+    readBy: []
+  })
+  RETURN m.id as id,
+         m.swarmId as swarmId,
+         m.fromAgentId as fromAgentId,
+         m.toAgentId as toAgentId,
+         m.category as category,
+         m.timestamp as timestamp,
+         m.expiresAt as expiresAt
+`;
+/**
+ * Read messages for an agent. Returns messages that are:
+ * 1. Addressed to this agent specifically, OR
+ * 2. Broadcast (toAgentId is null) to the same swarm
+ * AND not yet expired.
+ * Optionally filters to unread-only (not in readBy list).
+ */
+const READ_MESSAGES_QUERY = `
+  MATCH (m:SwarmMessage)
+  WHERE m.projectId = $projectId
+    AND m.swarmId = $swarmId
+    AND m.expiresAt > timestamp()
+    AND (m.toAgentId IS NULL OR m.toAgentId = $agentId)
+    AND ($unreadOnly = false OR NOT $agentId IN m.readBy)
+    AND ($categories IS NULL OR size($categories) = 0 OR m.category IN $categories)
+    AND ($fromAgentId IS NULL OR m.fromAgentId = $fromAgentId)
+  RETURN m.id as id,
+         m.swarmId as swarmId,
+         m.fromAgentId as fromAgentId,
+         m.toAgentId as toAgentId,
+         m.category as category,
+         m.content as content,
+         m.taskId as taskId,
+         m.filePaths as filePaths,
+         m.timestamp as timestamp,
+         m.expiresAt as expiresAt,
+         m.readBy as readBy,
+         NOT $agentId IN m.readBy as isUnread
+  ORDER BY m.timestamp DESC
+  LIMIT toInteger($limit)
+`;
+/**
+ * Acknowledge (mark as read) specific messages for an agent.
+ * Uses APOC to atomically add agentId to readBy array.
+ */
+const ACKNOWLEDGE_MESSAGES_QUERY = `
+  UNWIND $messageIds as msgId
+  MATCH (m:SwarmMessage {id: msgId, projectId: $projectId})
+  WHERE NOT $agentId IN m.readBy
+  SET m.readBy = m.readBy + $agentId
+  RETURN m.id as id, m.category as category
+`;
+/**
+ * Acknowledge ALL unread messages for an agent in a swarm.
+ */
+const ACKNOWLEDGE_ALL_QUERY = `
+  MATCH (m:SwarmMessage)
+  WHERE m.projectId = $projectId
+    AND m.swarmId = $swarmId
+    AND (m.toAgentId IS NULL OR m.toAgentId = $agentId)
+    AND NOT $agentId IN m.readBy
+    AND m.expiresAt > timestamp()
+  SET m.readBy = m.readBy + $agentId
+  RETURN count(m) as acknowledged
+`;
+/**
+ * Fetch pending messages for delivery during task claim.
+ * Returns unread messages addressed to or broadcast to the agent.
+ * Used internally by swarm_claim_task integration.
+ */
+export const PENDING_MESSAGES_FOR_AGENT_QUERY = `
+  MATCH (m:SwarmMessage)
+  WHERE m.projectId = $projectId
+    AND m.swarmId = $swarmId
+    AND m.expiresAt > timestamp()
+    AND (m.toAgentId IS NULL OR m.toAgentId = $agentId)
+    AND NOT $agentId IN m.readBy
+  RETURN m.id as id,
+         m.fromAgentId as fromAgentId,
+         m.category as category,
+         m.content as content,
+         m.taskId as taskId,
+         m.filePaths as filePaths,
+         m.timestamp as timestamp
+  ORDER BY
+    CASE m.category
+      WHEN 'alert' THEN 0
+      WHEN 'conflict' THEN 1
+      WHEN 'blocked' THEN 2
+      WHEN 'request' THEN 3
+      WHEN 'finding' THEN 4
+      WHEN 'handoff' THEN 5
+      ELSE 6
+    END,
+    m.timestamp DESC
+  LIMIT 10
+`;
+/**
+ * Auto-acknowledge messages that were delivered during claim.
+ */
+export const AUTO_ACKNOWLEDGE_QUERY = `
+  UNWIND $messageIds as msgId
+  MATCH (m:SwarmMessage {id: msgId})
+  WHERE NOT $agentId IN m.readBy
+  SET m.readBy = m.readBy + $agentId
+  RETURN count(m) as acknowledged
+`;
+/**
+ * Cleanup expired messages for a swarm.
+ */
+const CLEANUP_EXPIRED_QUERY = `
+  MATCH (m:SwarmMessage)
+  WHERE m.projectId = $projectId
+    AND ($swarmId IS NULL OR m.swarmId = $swarmId)
+    AND m.expiresAt < timestamp()
+  DELETE m
+  RETURN count(m) as cleaned
+`;
+// ============================================================================
+// TOOL CREATION
+// ============================================================================
+export const createSwarmMessageTool = (server) => {
+    server.registerTool(TOOL_NAMES.swarmMessage, {
+        title: TOOL_METADATA[TOOL_NAMES.swarmMessage].title,
+        description: TOOL_METADATA[TOOL_NAMES.swarmMessage].description,
+        inputSchema: {
+            projectId: z.string().describe('Project ID, name, or path'),
+            swarmId: z.string().describe('Swarm ID for scoping messages'),
+            agentId: z.string().describe('Your unique agent identifier'),
+            action: z
+                .enum(['send', 'read', 'acknowledge'])
+                .describe('Action: send (post message), read (get messages), acknowledge (mark as read)'),
+            // Send parameters
+            toAgentId: z
+                .string()
+                .optional()
+                .describe('Target agent ID. Omit for broadcast to all swarm agents.'),
+            category: z
+                .enum(MESSAGE_CATEGORY_KEYS)
+                .optional()
+                .describe('Message category: blocked (need help), conflict (resource clash), finding (important discovery), ' +
+                'request (direct ask), alert (urgent notification), handoff (context transfer)'),
+            content: z
+                .string()
+                .optional()
+                .describe('Message content (required for send action)'),
+            taskId: z
+                .string()
+                .optional()
+                .describe('Related task ID for context'),
+            filePaths: z
+                .array(z.string())
+                .optional()
+                .describe('File paths relevant to this message'),
+            ttlMs: z
+                .number()
+                .int()
+                .optional()
+                .describe(`Time-to-live in ms (default: ${MESSAGE_DEFAULT_TTL_MS / 3600000}h). Set 0 for swarm lifetime.`),
+            // Read parameters
+            unreadOnly: z
+                .boolean()
+                .optional()
+                .default(true)
+                .describe('Only return unread messages (default: true)'),
+            categories: z
+                .array(z.enum(MESSAGE_CATEGORY_KEYS))
+                .optional()
+                .describe('Filter by message categories'),
+            fromAgentId: z
+                .string()
+                .optional()
+                .describe('Filter messages from a specific agent'),
+            limit: z
+                .number()
+                .int()
+                .min(1)
+                .max(100)
+                .optional()
+                .default(20)
+                .describe('Maximum messages to return (default: 20)'),
+            // Acknowledge parameters
+            messageIds: z
+                .array(z.string())
+                .optional()
+                .describe('Specific message IDs to acknowledge. Omit to acknowledge all unread.'),
+            // Maintenance
+            cleanup: z
+                .boolean()
+                .optional()
+                .default(false)
+                .describe('Also clean up expired messages'),
+        },
+    }, async ({ projectId, swarmId, agentId, action, toAgentId, category, content, taskId, filePaths, ttlMs, unreadOnly = true, categories, fromAgentId, limit = 20, messageIds, cleanup = false, }) => {
+        const neo4jService = new Neo4jService();
+        const projectResult = await resolveProjectIdOrError(projectId, neo4jService);
+        if (!projectResult.success) {
+            await neo4jService.close();
+            return projectResult.error;
+        }
+        const resolvedProjectId = projectResult.projectId;
+        try {
+            // Optional cleanup of expired messages
+            let cleanedCount = 0;
+            if (cleanup) {
+                const cleanResult = await neo4jService.run(CLEANUP_EXPIRED_QUERY, {
+                    projectId: resolvedProjectId,
+                    swarmId,
+                });
+                cleanedCount = cleanResult[0]?.cleaned ?? 0;
+            }
+            // ── SEND ──────────────────────────────────────────────────────
+            if (action === 'send') {
+                if (!category) {
+                    return createErrorResponse('category is required for send action');
+                }
+                if (!content) {
+                    return createErrorResponse('content is required for send action');
+                }
+                const messageId = generateMessageId();
+                const effectiveTtl = ttlMs ?? MESSAGE_DEFAULT_TTL_MS;
+                const result = await neo4jService.run(SEND_MESSAGE_QUERY, {
+                    messageId,
+                    projectId: resolvedProjectId,
+                    swarmId,
+                    fromAgentId: agentId,
+                    toAgentId: toAgentId ?? null,
+                    category,
+                    content,
+                    taskId: taskId ?? null,
+                    filePaths: filePaths ?? [],
+                    ttlMs: effectiveTtl,
+                });
+                if (result.length === 0) {
+                    return createErrorResponse('Failed to create message');
+                }
+                const msg = result[0];
+                const ts = typeof msg.timestamp === 'object' && msg.timestamp?.toNumber ? msg.timestamp.toNumber() : msg.timestamp;
+                return createSuccessResponse(JSON.stringify({
+                    action: 'sent',
+                    message: {
+                        id: messageId,
+                        swarmId: msg.swarmId,
+                        from: msg.fromAgentId,
+                        to: msg.toAgentId ?? 'broadcast',
+                        category: msg.category,
+                        expiresIn: effectiveTtl > 0 ? `${Math.round(effectiveTtl / 60000)} minutes` : 'never',
+                    },
+                    ...(cleanedCount > 0 && { expiredCleaned: cleanedCount }),
+                }));
+            }
+            // ── READ ──────────────────────────────────────────────────────
+            if (action === 'read') {
+                const result = await neo4jService.run(READ_MESSAGES_QUERY, {
+                    projectId: resolvedProjectId,
+                    swarmId,
+                    agentId,
+                    unreadOnly,
+                    categories: categories ?? null,
+                    fromAgentId: fromAgentId ?? null,
+                    limit: Math.floor(limit),
+                });
+                const messages = result.map((m) => {
+                    const ts = typeof m.timestamp === 'object' && m.timestamp?.toNumber ? m.timestamp.toNumber() : m.timestamp;
+                    return {
+                        id: m.id,
+                        from: m.fromAgentId,
+                        to: m.toAgentId ?? 'broadcast',
+                        category: m.category,
+                        content: m.content,
+                        taskId: m.taskId ?? undefined,
+                        filePaths: m.filePaths?.length > 0 ? m.filePaths : undefined,
+                        isUnread: m.isUnread,
+                        age: ts ? `${Math.round((Date.now() - ts) / 1000)}s ago` : null,
+                    };
+                });
+                return createSuccessResponse(JSON.stringify({
+                    action: 'read',
+                    swarmId,
+                    forAgent: agentId,
+                    count: messages.length,
+                    messages,
+                    ...(cleanedCount > 0 && { expiredCleaned: cleanedCount }),
+                }));
+            }
+            // ── ACKNOWLEDGE ───────────────────────────────────────────────
+            if (action === 'acknowledge') {
+                if (messageIds && messageIds.length > 0) {
+                    // Acknowledge specific messages
+                    const result = await neo4jService.run(ACKNOWLEDGE_MESSAGES_QUERY, {
+                        messageIds,
+                        projectId: resolvedProjectId,
+                        agentId,
+                    });
+                    return createSuccessResponse(JSON.stringify({
+                        action: 'acknowledged',
+                        count: result.length,
+                        messageIds: result.map((r) => r.id),
+                        ...(cleanedCount > 0 && { expiredCleaned: cleanedCount }),
+                    }));
+                }
+                else {
+                    // Acknowledge all unread
+                    const result = await neo4jService.run(ACKNOWLEDGE_ALL_QUERY, {
+                        projectId: resolvedProjectId,
+                        swarmId,
+                        agentId,
+                    });
+                    const count = typeof result[0]?.acknowledged === 'object'
+                        ? result[0].acknowledged.toNumber()
+                        : result[0]?.acknowledged ?? 0;
+                    return createSuccessResponse(JSON.stringify({
+                        action: 'acknowledged_all',
+                        count,
+                        ...(cleanedCount > 0 && { expiredCleaned: cleanedCount }),
+                    }));
+                }
+            }
+            return createErrorResponse(`Unknown action: ${action}`);
+        }
+        catch (error) {
+            await debugLog('Swarm message error', { error: String(error) });
+            return createErrorResponse(error instanceof Error ? error : String(error));
+        }
+        finally {
+            await neo4jService.close();
+        }
+    });
+};

package/dist/storage/neo4j/neo4j.service.js

@@ -136,19 +136,19 @@ export const QUERIES = {
     RETURN labels(n)[0] as nodeType, count(*) as count
     ORDER BY count DESC
   `,
-    CREATE_EMBEDDED_VECTOR_INDEX: `
+    CREATE_EMBEDDED_VECTOR_INDEX: (dimensions) => `
   CREATE VECTOR INDEX embedded_nodes_idx IF NOT EXISTS
   FOR (n:Embedded) ON (n.embedding)
   OPTIONS {indexConfig: {
-    \`vector.dimensions\`: 3072,
+    \`vector.dimensions\`: ${dimensions},
     \`vector.similarity_function\`: 'cosine'
   }}
 `,
-    CREATE_SESSION_NOTES_VECTOR_INDEX: `
+    CREATE_SESSION_NOTES_VECTOR_INDEX: (dimensions) => `
   CREATE VECTOR INDEX session_notes_idx IF NOT EXISTS
   FOR (n:SessionNote) ON (n.embedding)
   OPTIONS {indexConfig: {
-    \`vector.dimensions\`: 3072,
+    \`vector.dimensions\`: ${dimensions},
     \`vector.similarity_function\`: 'cosine'
   }}
 `,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "code-graph-context",
-  "version": "2.8.0",
+  "version": "2.10.0",
   "description": "MCP server that builds code graphs to provide rich context to LLMs",
   "type": "module",
   "homepage": "https://github.com/drewdrewH/code-graph-context#readme",
@@ -34,6 +34,8 @@
   },
   "files": [
     "dist/**/*",
+    "sidecar/embedding_server.py",
+    "sidecar/requirements.txt",
     "README.md",
     "LICENSE",
     ".env.example"

package/sidecar/embedding_server.py

@@ -0,0 +1,147 @@
+"""
+Local embedding server for code-graph-context.
+Uses Qodo-Embed-1-1.5B for high-quality code embeddings without OpenAI dependency.
+Runs as a sidecar process managed by the Node.js MCP server.
+"""
+
+import gc
+import os
+import sys
+import signal
+import logging
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    stream=sys.stderr,
+)
+logger = logging.getLogger("embedding-sidecar")
+
+app = FastAPI(title="code-graph-context embedding sidecar")
+
+model = None
+model_name = os.environ.get("EMBEDDING_MODEL", "Qodo/Qodo-Embed-1-1.5B")
+
+
+class EmbedRequest(BaseModel):
+    texts: list[str]
+    batch_size: int = 8
+
+
+class EmbedResponse(BaseModel):
+    embeddings: list[list[float]]
+    dimensions: int
+    model: str
+
+
+@app.on_event("startup")
+def load_model():
+    global model
+    try:
+        import torch
+        from sentence_transformers import SentenceTransformer
+
+        device = "mps" if torch.backends.mps.is_available() else "cpu"
+        logger.info(f"Loading {model_name} on {device}...")
+        model = SentenceTransformer(model_name, device=device)
+
+        # Warm up with a test embedding
+        test = model.encode(["warmup"], show_progress_bar=False)
+        dims = len(test[0])
+        logger.info(f"Model loaded: {dims} dimensions, device={device}")
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+
+@app.get("/health")
+def health():
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    sample = model.encode(["dim_check"], show_progress_bar=False)
+    return {
+        "status": "ok",
+        "model": model_name,
+        "dimensions": len(sample[0]),
+        "device": str(model.device),
+    }
+
+
+@app.post("/embed", response_model=EmbedResponse)
+async def embed(req: EmbedRequest):
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    if not req.texts:
+        return EmbedResponse(embeddings=[], dimensions=0, model=model_name)
+
+    try:
+        embeddings = _encode_with_oom_fallback(req.texts, req.batch_size)
+        dims = len(embeddings[0])
+        return EmbedResponse(
+            embeddings=embeddings,
+            dimensions=dims,
+            model=model_name,
+        )
+    except Exception as e:
+        logger.error(f"Embedding error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def _encode_with_oom_fallback(texts: list[str], batch_size: int) -> list[list[float]]:
+    """
+    Encode texts, falling back to CPU if MPS runs out of memory.
+    Also retries with smaller batch sizes before giving up.
+    """
+    import torch
+
+    try:
+        result = model.encode(
+            texts,
+            batch_size=batch_size,
+            show_progress_bar=False,
+            normalize_embeddings=True,
+        )
+        return result.tolist()
+    except (torch.mps.OutOfMemoryError, RuntimeError) as e:
+        if "out of memory" not in str(e).lower():
+            raise
+
+        logger.warning(f"MPS OOM with batch_size={batch_size}, len={len(texts)}. Falling back to CPU.")
+
+        # Free MPS memory
+        if hasattr(torch.mps, "empty_cache"):
+            torch.mps.empty_cache()
+        gc.collect()
+
+        # Fall back to CPU for this request
+        original_device = model.device
+        model.to("cpu")
+
+        try:
+            # Use smaller batches on CPU
+            cpu_batch = min(batch_size, 4)
+            result = model.encode(
+                texts,
+                batch_size=cpu_batch,
+                show_progress_bar=False,
+                normalize_embeddings=True,
+            )
+            return result.tolist()
+        finally:
+            # Move back to MPS for future requests
+            try:
+                model.to(original_device)
+            except Exception:
+                logger.warning("Could not move model back to MPS, staying on CPU")
+
+
+def handle_signal(sig, _frame):
+    logger.info(f"Received signal {sig}, shutting down")
+    sys.exit(0)
+
+
+signal.signal(signal.SIGTERM, handle_signal)

package/sidecar/requirements.txt

@@ -0,0 +1,5 @@
+fastapi>=0.104.0
+uvicorn>=0.24.0
+sentence-transformers>=3.0.0
+torch>=2.0.0
+pydantic>=2.0.0