import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { McpDependencies } from '../server/mcp-server.js';
import { memoryUpsertInput } from './schemas.js';
import { embed } from '../services/ollama-client.js';
import { chunkText } from '../services/chunker.js';
import { contentHash } from '../services/dedup.js';
import { auditToolCall } from '../middleware/audit.js';
import { logger } from '../utils/logger.js';
export function registerMemoryUpsert(server: McpServer, deps: McpDependencies): void {
server.registerTool('memory_upsert', {
description:
'Store or update a document in memory. Content is chunked, embedded, and indexed for future semantic search.',
inputSchema: memoryUpsertInput,
}, async (args, extra) => {
const start = Date.now();
const { document_id, content, scope = 'default', tags = [], source } = args;
try {
// Check if content has changed via hash
const hash = contentHash(content);
const existing = deps.db
.prepare('SELECT content_hash FROM documents WHERE id = ?')
.get(document_id) as { content_hash: string } | undefined;
if (existing?.content_hash === hash) {
const output = JSON.stringify({
document_id,
status: 'unchanged',
message: 'Content hash matches existing document, skipping.',
});
auditToolCall(deps.db, 'memory_upsert', extra.sessionId, document_id, 'unchanged', Date.now() - start);
return { content: [{ type: 'text' as const, text: output }] };
}
// Chunk the content
const chunks = chunkText(content);
logger.info({ document_id, chunks: chunks.length }, 'Chunking document');
// Embed all chunks in batch
const embeddings = await embed(chunks);
// Upsert to vector store
const chunksWithEmbeddings = chunks.map((text, i) => ({
content: text,
embedding: embeddings[i],
}));
await deps.vectorStore.upsertDocument(document_id, chunksWithEmbeddings, {
scope,
tags,
source,
});
// Update content hash
deps.db
.prepare('UPDATE documents SET content_hash = ? WHERE id = ?')
.run(hash, document_id);
const output = JSON.stringify({
document_id,
status: 'upserted',
chunks_count: chunks.length,
scope,
});
auditToolCall(deps.db, 'memory_upsert', extra.sessionId, document_id, `${chunks.length} chunks`, Date.now() - start);
return { content: [{ type: 'text' as const, text: output }] };
} catch (err) {
const msg = err instanceof Error ? err.message : String(err);
logger.error({ err, document_id }, 'memory.upsert failed');
auditToolCall(deps.db, 'memory_upsert', extra.sessionId, document_id, `error: ${msg}`, Date.now() - start);
return { content: [{ type: 'text' as const, text: `Error: ${msg}` }], isError: true };
}
});
}
