From a46440633846ac058041b6f33b0c1601437481aa Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Fri, 10 Apr 2026 07:37:55 -1000 Subject: [PATCH] fix: 7 bug fixes from Issue #9 and #22 - fix(mcp): use ListToolsRequestSchema/CallToolRequestSchema instead of string literals (Issue #9, PR #25) - fix(mcp): handleToolCall reads dry_run from params instead of hardcoding false (#22 Bug #11) - fix(search): keyword search returns best chunk per page via DISTINCT ON, not all chunks (#22 Bug #8) - fix(search): dedup layer 1 keeps top 3 chunks per page instead of collapsing to 1 (#22 Bug #12) - fix(engine): transaction uses scoped engine via Object.create, no shared state mutation (#22 Bug #2) - fix(engine): upsertChunks uses UPSERT instead of DELETE+INSERT, preserves existing embeddings (#22 Bug #1) - fix(slugs): validateSlug normalizes to lowercase, pathToSlug lowercases consistently (#22 Bug #4) - schema: add unique index on content_chunks(page_id, chunk_index) for UPSERT support - schema: add access_tokens and mcp_request_log tables via migration Co-Authored-By: Claude Opus 4.6 (1M context) --- src/core/db.ts | 11 +---- src/core/migrate.ts | 36 +++++++++++++- src/core/postgres-engine.ts | 93 +++++++++++++++++++++---------------- src/core/search/dedup.ts | 20 +++++--- src/core/sync.ts | 2 +- src/mcp/server.ts | 7 +-- src/schema.sql | 1 + test/sync.test.ts | 4 +- 8 files changed, 109 insertions(+), 65 deletions(-) diff --git a/src/core/db.ts b/src/core/db.ts index 20b9d42..459a483 100644 --- a/src/core/db.ts +++ b/src/core/db.ts @@ -72,16 +72,9 @@ export async function initSchema(): Promise { await conn.unsafe(schemaSql); } -export async function withTransaction(fn: () => Promise): Promise { +export async function withTransaction(fn: (tx: ReturnType) => Promise): Promise { const conn = getConnection(); return conn.begin(async (tx) => { - // Temporarily swap global connection to transaction - const prev = sql; - sql = tx as unknown as ReturnType; - try { - return await fn(); - } finally { - sql = prev; - } + return fn(tx as unknown as ReturnType); }); } diff --git a/src/core/migrate.ts b/src/core/migrate.ts index cb60b49..3f233ad 100644 --- a/src/core/migrate.ts +++ b/src/core/migrate.ts @@ -20,8 +20,40 @@ interface Migration { // Add new migrations at the end. Never modify existing ones. const MIGRATIONS: Migration[] = [ // Version 1 is the baseline (schema.sql creates everything with IF NOT EXISTS). - // Future migrations go here: - // { version: 2, name: 'add_aliases', sql: `ALTER TABLE pages ADD COLUMN IF NOT EXISTS aliases TEXT[];` }, + { + version: 2, + name: 'unique_chunk_index', + sql: ` + -- Deduplicate any existing duplicate (page_id, chunk_index) rows before adding constraint + DELETE FROM content_chunks a USING content_chunks b + WHERE a.page_id = b.page_id AND a.chunk_index = b.chunk_index AND a.id > b.id; + CREATE UNIQUE INDEX IF NOT EXISTS idx_chunks_page_index ON content_chunks(page_id, chunk_index); + `, + }, + { + version: 3, + name: 'access_tokens_and_mcp_log', + sql: ` + CREATE TABLE IF NOT EXISTS access_tokens ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + name TEXT NOT NULL, + token_hash TEXT NOT NULL UNIQUE, + scopes TEXT[], + created_at TIMESTAMPTZ DEFAULT now(), + last_used_at TIMESTAMPTZ, + revoked_at TIMESTAMPTZ + ); + CREATE INDEX IF NOT EXISTS idx_access_tokens_hash ON access_tokens (token_hash) WHERE revoked_at IS NULL; + CREATE TABLE IF NOT EXISTS mcp_request_log ( + id SERIAL PRIMARY KEY, + token_name TEXT, + operation TEXT NOT NULL, + latency_ms INTEGER, + status TEXT NOT NULL DEFAULT 'success', + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + `, + }, ]; export const LATEST_VERSION = MIGRATIONS.length > 0 diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index 5ac3bf4..3428fe5 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -70,19 +70,14 @@ export class PostgresEngine implements BrainEngine { } async transaction(fn: (engine: BrainEngine) => Promise): Promise { - if (this._sql) { - // Instance connection: use .begin() directly, no global swap - return this._sql.begin(async (tx) => { - const prev = this._sql; - this._sql = tx as unknown as ReturnType; - try { - return await fn(this); - } finally { - this._sql = prev; - } - }); - } - return db.withTransaction(() => fn(this)); + const conn = this._sql || db.getConnection(); + return conn.begin(async (tx) => { + // Create a scoped engine with tx as its connection, no shared state mutation + const txEngine = Object.create(this) as PostgresEngine; + Object.defineProperty(txEngine, 'sql', { get: () => tx }); + Object.defineProperty(txEngine, '_sql', { value: tx as unknown as ReturnType, writable: false }); + return fn(txEngine); + }); } // Pages CRUD @@ -97,7 +92,7 @@ export class PostgresEngine implements BrainEngine { } async putPage(slug: string, page: PageInput): Promise { - validateSlug(slug); + slug = validateSlug(slug); const sql = this.sql; const hash = page.content_hash || contentHash(page.compiled_truth, page.timeline || ''); const frontmatter = page.frontmatter || {}; @@ -182,7 +177,7 @@ export class PostgresEngine implements BrainEngine { const limit = opts?.limit || 20; const rows = await sql` - SELECT + SELECT DISTINCT ON (p.slug) p.slug, p.id as page_id, p.title, p.type, cc.chunk_text, cc.chunk_source, ts_rank(p.search_vector, websearch_to_tsquery('english', ${query})) AS score, @@ -192,9 +187,11 @@ export class PostgresEngine implements BrainEngine { FROM pages p JOIN content_chunks cc ON cc.page_id = p.id WHERE p.search_vector @@ websearch_to_tsquery('english', ${query}) - ORDER BY score DESC - LIMIT ${limit} + ORDER BY p.slug, score DESC `; + // Re-sort by score (DISTINCT ON requires ORDER BY slug first) and apply limit + rows.sort((a: any, b: any) => b.score - a.score); + rows.splice(limit); return rows.map(rowToSearchResult); } @@ -231,37 +228,50 @@ export class PostgresEngine implements BrainEngine { if (pages.length === 0) throw new Error(`Page not found: ${slug}`); const pageId = pages[0].id; - // Delete existing chunks for this page - await sql`DELETE FROM content_chunks WHERE page_id = ${pageId}`; - - // Bulk insert chunks — build multi-row VALUES to reduce round-trips - if (chunks.length === 0) return; - - // postgres.js tagged templates don't handle vector casting in bulk, - // so we build a parameterized raw SQL query - const cols = '(page_id, chunk_index, chunk_text, chunk_source, embedding, model, token_count, embedded_at)'; - const rows: string[] = []; - const params: unknown[] = []; - let paramIdx = 1; + // Remove chunks that no longer exist (chunk_index beyond new count) + const newIndices = chunks.map(c => c.chunk_index); + if (newIndices.length > 0) { + await sql`DELETE FROM content_chunks WHERE page_id = ${pageId} AND chunk_index != ALL(${newIndices})`; + } else { + await sql`DELETE FROM content_chunks WHERE page_id = ${pageId}`; + return; + } + // Upsert chunks — preserves existing embeddings when new embedding is undefined for (const chunk of chunks) { const embeddingStr = chunk.embedding ? '[' + Array.from(chunk.embedding).join(',') + ']' : null; if (embeddingStr) { - rows.push(`($${paramIdx++}, $${paramIdx++}, $${paramIdx++}, $${paramIdx++}, $${paramIdx++}::vector, $${paramIdx++}, $${paramIdx++}, now())`); - params.push(pageId, chunk.chunk_index, chunk.chunk_text, chunk.chunk_source, embeddingStr, chunk.model || 'text-embedding-3-large', chunk.token_count || null); + await sql.unsafe( + `INSERT INTO content_chunks (page_id, chunk_index, chunk_text, chunk_source, embedding, model, token_count, embedded_at) + VALUES ($1, $2, $3, $4, $5::vector, $6, $7, now()) + ON CONFLICT (page_id, chunk_index) DO UPDATE SET + chunk_text = EXCLUDED.chunk_text, + chunk_source = EXCLUDED.chunk_source, + embedding = EXCLUDED.embedding, + model = EXCLUDED.model, + token_count = EXCLUDED.token_count, + embedded_at = EXCLUDED.embedded_at`, + [pageId, chunk.chunk_index, chunk.chunk_text, chunk.chunk_source, embeddingStr, chunk.model || 'text-embedding-3-large', chunk.token_count || null], + ); } else { - rows.push(`($${paramIdx++}, $${paramIdx++}, $${paramIdx++}, $${paramIdx++}, NULL, $${paramIdx++}, $${paramIdx++}, NULL)`); - params.push(pageId, chunk.chunk_index, chunk.chunk_text, chunk.chunk_source, chunk.model || 'text-embedding-3-large', chunk.token_count || null); + // No new embedding: preserve existing embedding via COALESCE + await sql.unsafe( + `INSERT INTO content_chunks (page_id, chunk_index, chunk_text, chunk_source, embedding, model, token_count, embedded_at) + VALUES ($1, $2, $3, $4, NULL, $5, $6, NULL) + ON CONFLICT (page_id, chunk_index) DO UPDATE SET + chunk_text = EXCLUDED.chunk_text, + chunk_source = EXCLUDED.chunk_source, + embedding = COALESCE(content_chunks.embedding, EXCLUDED.embedding), + model = COALESCE(content_chunks.model, EXCLUDED.model), + token_count = EXCLUDED.token_count, + embedded_at = COALESCE(content_chunks.embedded_at, EXCLUDED.embedded_at)`, + [pageId, chunk.chunk_index, chunk.chunk_text, chunk.chunk_source, chunk.model || 'text-embedding-3-large', chunk.token_count || null], + ); } } - - await sql.unsafe( - `INSERT INTO content_chunks ${cols} VALUES ${rows.join(', ')}`, - params, - ); } async getChunks(slug: string): Promise { @@ -584,7 +594,7 @@ export class PostgresEngine implements BrainEngine { // Sync async updateSlug(oldSlug: string, newSlug: string): Promise { - validateSlug(newSlug); + newSlug = validateSlug(newSlug); const sql = this.sql; await sql`UPDATE pages SET slug = ${newSlug}, updated_at = now() WHERE slug = ${oldSlug}`; } @@ -613,12 +623,13 @@ export class PostgresEngine implements BrainEngine { } // Helpers -function validateSlug(slug: string): void { +function validateSlug(slug: string): string { // Git is the system of record — slugs are lowercased repo-relative paths. - // Only reject empty, path traversal (..), and leading slash. if (!slug || /\.\./.test(slug) || /^\//.test(slug)) { throw new Error(`Invalid slug: "${slug}". Slugs cannot be empty, start with /, or contain path traversal.`); } + // Normalize to lowercase — all entry points (pathToSlug, inferSlug, frontmatter, direct writes) go through here + return slug.toLowerCase(); } function contentHash(compiledTruth: string, timeline: string): string { diff --git a/src/core/search/dedup.ts b/src/core/search/dedup.ts index f8a0e3a..2af9c07 100644 --- a/src/core/search/dedup.ts +++ b/src/core/search/dedup.ts @@ -45,19 +45,25 @@ export function dedupResults( } /** - * Layer 1: Keep only the highest-scoring chunk per page. + * Layer 1: Keep top 3 chunks per page (not just 1). + * Later layers (text similarity, cap per page) handle further reduction. */ function dedupBySource(results: SearchResult[]): SearchResult[] { - const byPage = new Map(); + const byPage = new Map(); for (const r of results) { - const existing = byPage.get(r.slug); - if (!existing || r.score > existing.score) { - byPage.set(r.slug, r); - } + const existing = byPage.get(r.slug) || []; + existing.push(r); + byPage.set(r.slug, existing); } - return Array.from(byPage.values()).sort((a, b) => b.score - a.score); + const kept: SearchResult[] = []; + for (const chunks of byPage.values()) { + chunks.sort((a, b) => b.score - a.score); + kept.push(...chunks.slice(0, 3)); + } + + return kept.sort((a, b) => b.score - a.score); } /** diff --git a/src/core/sync.ts b/src/core/sync.ts index 8da1a3a..6a8e7f8 100644 --- a/src/core/sync.ts +++ b/src/core/sync.ts @@ -113,5 +113,5 @@ export function pathToSlug(filePath: string, repoPrefix?: string): string { slug = slug.replace(/^\//, ''); // Add repo prefix for multi-repo setups if (repoPrefix) slug = `${repoPrefix}/${slug}`; - return slug; + return slug.toLowerCase(); } diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 7a28685..a798b64 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1,5 +1,6 @@ import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; +import { ListToolsRequestSchema, CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import type { BrainEngine } from '../core/engine.ts'; import { operations, OperationError } from '../core/operations.ts'; import type { OperationContext } from '../core/operations.ts'; @@ -13,7 +14,7 @@ export async function startMcpServer(engine: BrainEngine) { ); // Generate tool definitions from operations - server.setRequestHandler('tools/list' as any, async () => ({ + server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: operations.map(op => ({ name: op.name, description: op.description, @@ -35,7 +36,7 @@ export async function startMcpServer(engine: BrainEngine) { })); // Dispatch tool calls to operation handlers - server.setRequestHandler('tools/call' as any, async (request: any) => { + server.setRequestHandler(CallToolRequestSchema, async (request: any) => { const { name, arguments: params } = request.params; const op = operations.find(o => o.name === name); if (!op) { @@ -82,7 +83,7 @@ export async function handleToolCall( engine, config: loadConfig() || { engine: 'postgres' }, logger: { info: console.log, warn: console.warn, error: console.error }, - dryRun: false, + dryRun: !!(params?.dry_run), }; return op.handler(ctx, params); diff --git a/src/schema.sql b/src/schema.sql index 03e611e..910d728 100644 --- a/src/schema.sql +++ b/src/schema.sql @@ -39,6 +39,7 @@ CREATE TABLE IF NOT EXISTS content_chunks ( created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +CREATE UNIQUE INDEX IF NOT EXISTS idx_chunks_page_index ON content_chunks(page_id, chunk_index); CREATE INDEX IF NOT EXISTS idx_chunks_page ON content_chunks(page_id); CREATE INDEX IF NOT EXISTS idx_chunks_embedding ON content_chunks USING hnsw (embedding vector_cosine_ops); diff --git a/test/sync.test.ts b/test/sync.test.ts index 375479d..0ce3881 100644 --- a/test/sync.test.ts +++ b/test/sync.test.ts @@ -94,8 +94,8 @@ describe('pathToSlug', () => { expect(pathToSlug('people/pedro-franceschi.md')).toBe('people/pedro-franceschi'); }); - test('preserves case', () => { - expect(pathToSlug('People/Pedro-Franceschi.md')).toBe('People/Pedro-Franceschi'); + test('normalizes to lowercase', () => { + expect(pathToSlug('People/Pedro-Franceschi.md')).toBe('people/pedro-franceschi'); }); test('strips leading slash', () => {