triton6564685/gbrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 7 bug fixes from Issue #9 and #22

Browse Source 
- 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) <noreply@anthropic.com>
This commit is contained in:
Garry Tan
2026-04-10 07:37:55 -10:00
parent c68a4ccbbb
commit a464406338
8 changed files with 109 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/core/db.ts
View File

@@ -72,16 +72,9 @@ export async function initSchema(): Promise<void> {
await conn.unsafe(schemaSql);
}
export async function withTransaction<T>(fn: () => Promise<T>): Promise<T> {
export async function withTransaction<T>(fn: (tx: ReturnType<typeof postgres>) => Promise<T>): Promise<T> {
const conn = getConnection();
return conn.begin(async (tx) => {
// Temporarily swap global connection to transaction
const prev = sql;
sql = tx as unknown as ReturnType<typeof postgres>;
try {
return await fn();
} finally {
sql = prev;
}
return fn(tx as unknown as ReturnType<typeof postgres>);
});
}

36
src/core/migrate.ts
View File

@@ -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

91
src/core/postgres-engine.ts
View File

@@ -70,20 +70,15 @@ export class PostgresEngine implements BrainEngine {
}
async transaction<T>(fn: (engine: BrainEngine) => Promise<T>): Promise<T> {
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<typeof postgres>;
try {
return await fn(this);
} finally {
this._sql = prev;
}
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<typeof postgres>, writable: false });
return fn(txEngine);
});
}
return db.withTransaction(() => fn(this));
}
// Pages CRUD
async getPage(slug: string): Promise<Page | null> {
@@ -97,7 +92,7 @@ export class PostgresEngine implements BrainEngine {
}
async putPage(slug: string, page: PageInput): Promise<Page> {
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
// 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;
}
// 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;
// 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);
} 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);
}
}
await sql.unsafe(
`INSERT INTO content_chunks ${cols} VALUES ${rows.join(', ')}`,
params,
`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 {
// 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],
);
}
}
}
async getChunks(slug: string): Promise<Chunk[]> {
@@ -584,7 +594,7 @@ export class PostgresEngine implements BrainEngine {
// Sync
async updateSlug(oldSlug: string, newSlug: string): Promise<void> {
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 {

20
src/core/search/dedup.ts
View File

@@ -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<string, SearchResult>();
const byPage = new Map<string, SearchResult[]>();
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);
}
/**

2
src/core/sync.ts
View File

@@ -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();
}

7
src/mcp/server.ts
View File

@@ -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);

1
src/schema.sql
View File

@@ -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);

4
test/sync.test.ts
View File

@@ -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', () => {