From 010847c02016450c77ccfc841d3b233471eb96b5 Mon Sep 17 00:00:00 2001 From: Time Attakc <89218912+time-attack@users.noreply.github.com> Date: Mon, 13 Jul 2026 00:09:27 -0700 Subject: [PATCH] fix(think): enforce source scope across gather (#2200) (#2739) Carry the caller's scalar or federated source scope through every think gather stream: hybrid page retrieval, takes keyword/vector retrieval, and graph traversal. Adds source predicates to the takes retrieval methods in both engines. Part of #2200 (the think slice; #2200 stays open as the tracking issue for the remaining by-slug read ops). --- src/core/operations.ts | 18 +++- src/core/pglite-engine.ts | 24 ++++- src/core/postgres-engine.ts | 12 +++ src/core/think/gather.ts | 13 ++- src/core/think/index.ts | 2 + .../e2e/think-source-isolation-pglite.test.ts | 92 +++++++++++++++++++ test/source-scope-resolver.test.ts | 15 +++ 7 files changed, 168 insertions(+), 8 deletions(-) create mode 100644 test/e2e/think-source-isolation-pglite.test.ts diff --git a/src/core/operations.ts b/src/core/operations.ts index e8d28913c..41f376b6f 100644 --- a/src/core/operations.ts +++ b/src/core/operations.ts @@ -425,6 +425,19 @@ export function sourceScopeOpts(ctx: OperationContext): { sourceId?: string; sou return {}; } +/** Map the operation-layer scope names onto runThink's public options. */ +export function thinkSourceScopeOpts(ctx: OperationContext): { + sourceId?: string; + allowedSources?: string[]; +} { + const scope = sourceScopeOpts(ctx); + return scope.sourceIds !== undefined + ? { allowedSources: scope.sourceIds } + : scope.sourceId !== undefined + ? { sourceId: scope.sourceId } + : {}; +} + /** * #2200: source scope for the LINK read ops (get_links / get_backlinks). A link * row references three pages (from, to, origin); the engine's federated @@ -1842,7 +1855,7 @@ const think: Operation = { // present) OR the scalar; we pass both through to runThink which // forwards to findTrajectory. CLI callers don't go through this op // and get default scope + remote=false from runThink's CLI path. - const scope = sourceScopeOpts(ctx); + const thinkScope = thinkSourceScopeOpts(ctx); const { runThink, persistSynthesis } = await import('./think/index.ts'); const result = await runThink(ctx.engine, { question: String(p.question), @@ -1859,8 +1872,7 @@ const think: Operation = { since: p.since ? String(p.since) : undefined, until: p.until ? String(p.until) : undefined, takesHoldersAllowList: ctx.takesHoldersAllowList, - ...(scope.sourceId !== undefined ? { sourceId: scope.sourceId } : {}), - ...(scope.sourceIds !== undefined ? { allowedSources: scope.sourceIds } : {}), + ...thinkScope, remote: ctx.remote === true, }); diff --git a/src/core/pglite-engine.ts b/src/core/pglite-engine.ts index df686ac2f..c5bf8f1f6 100644 --- a/src/core/pglite-engine.ts +++ b/src/core/pglite-engine.ts @@ -4603,7 +4603,7 @@ export class PGLiteEngine implements BrainEngine { async searchTakes( query: string, - opts: { limit?: number; takesHoldersAllowList?: string[] } = {}, + opts: SearchOpts & { takesHoldersAllowList?: string[] } = {}, ): Promise { const limit = clampSearchLimit(opts.limit, 30, 100); const { rows } = await this.db.query( @@ -4615,16 +4615,24 @@ export class PGLiteEngine implements BrainEngine { WHERE t.active AND t.claim % $1 AND ($2::text[] IS NULL OR t.holder = ANY($2::text[])) + AND ($4::text[] IS NULL OR p.source_id = ANY($4::text[])) + AND ($5::text IS NULL OR p.source_id = $5::text) ORDER BY score DESC, t.weight DESC LIMIT $3`, - [query, opts.takesHoldersAllowList ?? null, limit] + [ + query, + opts.takesHoldersAllowList ?? null, + limit, + opts.sourceIds && opts.sourceIds.length > 0 ? opts.sourceIds : null, + opts.sourceIds && opts.sourceIds.length > 0 ? null : (opts.sourceId ?? null), + ] ); return rows as unknown as TakeHit[]; } async searchTakesVector( embedding: Float32Array, - opts: { limit?: number; takesHoldersAllowList?: string[] } = {}, + opts: SearchOpts & { takesHoldersAllowList?: string[] } = {}, ): Promise { const limit = clampSearchLimit(opts.limit, 30, 100); const vec = `[${Array.from(embedding).join(',')}]`; @@ -4637,9 +4645,17 @@ export class PGLiteEngine implements BrainEngine { WHERE t.active AND t.embedding IS NOT NULL AND ($2::text[] IS NULL OR t.holder = ANY($2::text[])) + AND ($4::text[] IS NULL OR p.source_id = ANY($4::text[])) + AND ($5::text IS NULL OR p.source_id = $5::text) ORDER BY t.embedding <=> $1::vector LIMIT $3`, - [vec, opts.takesHoldersAllowList ?? null, limit] + [ + vec, + opts.takesHoldersAllowList ?? null, + limit, + opts.sourceIds && opts.sourceIds.length > 0 ? opts.sourceIds : null, + opts.sourceIds && opts.sourceIds.length > 0 ? null : (opts.sourceId ?? null), + ] ); return rows as unknown as TakeHit[]; } diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index ce0d06e5b..1ed51ce9a 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -4615,6 +4615,11 @@ export class PostgresEngine implements BrainEngine { async searchTakes(query: string, opts: SearchOpts & { takesHoldersAllowList?: string[] } = {}): Promise { const sql = this.sql; const limit = clampSearchLimit(opts.limit, 30, 100); + const sourceFilter = opts.sourceIds && opts.sourceIds.length > 0 + ? sql`AND p.source_id = ANY(${opts.sourceIds}::text[])` + : opts.sourceId + ? sql`AND p.source_id = ${opts.sourceId}` + : sql``; const rows = await sql` SELECT t.id AS take_id, t.page_id, p.slug AS page_slug, t.row_num, t.claim, t.kind, t.holder, t.weight, @@ -4627,6 +4632,7 @@ export class PostgresEngine implements BrainEngine { ${opts.takesHoldersAllowList ?? null}::text[] IS NULL OR t.holder = ANY(${opts.takesHoldersAllowList ?? null}::text[]) ) + ${sourceFilter} ORDER BY score DESC, t.weight DESC LIMIT ${limit} `; @@ -4640,6 +4646,11 @@ export class PostgresEngine implements BrainEngine { const sql = this.sql; const limit = clampSearchLimit(opts.limit, 30, 100); const vec = `[${Array.from(embedding).join(',')}]`; + const sourceFilter = opts.sourceIds && opts.sourceIds.length > 0 + ? sql`AND p.source_id = ANY(${opts.sourceIds}::text[])` + : opts.sourceId + ? sql`AND p.source_id = ${opts.sourceId}` + : sql``; const rows = await sql` SELECT t.id AS take_id, t.page_id, p.slug AS page_slug, t.row_num, t.claim, t.kind, t.holder, t.weight, @@ -4652,6 +4663,7 @@ export class PostgresEngine implements BrainEngine { ${opts.takesHoldersAllowList ?? null}::text[] IS NULL OR t.holder = ANY(${opts.takesHoldersAllowList ?? null}::text[]) ) + ${sourceFilter} ORDER BY t.embedding <=> ${vec}::vector LIMIT ${limit} `; diff --git a/src/core/think/gather.ts b/src/core/think/gather.ts index 11fcc3d1b..881df4c39 100644 --- a/src/core/think/gather.ts +++ b/src/core/think/gather.ts @@ -34,6 +34,9 @@ export interface ThinkGatherOpts { questionEmbedding?: Float32Array; /** When set, MCP-bound calls forward this allow-list to takes_search. Local CLI leaves unset. */ takesHoldersAllowList?: string[]; + /** Source scope inherited from the caller. Federated array wins over scalar. */ + sourceId?: string; + sourceIds?: string[]; } export interface ThinkGatherResult { @@ -101,6 +104,11 @@ export async function runGather( const gatherLimit = opts.gatherLimit ?? 40; const takesLimit = opts.takesLimit ?? 30; const graphDepth = opts.graphDepth ?? 2; + const sourceScope = opts.sourceIds && opts.sourceIds.length > 0 + ? { sourceIds: opts.sourceIds } + : opts.sourceId + ? { sourceId: opts.sourceId } + : {}; // Sanitize the question for any path that includes it in an LLM prompt. // (Direct DB search is fine — those are parameterized queries.) @@ -110,6 +118,7 @@ export async function runGather( const pagesPromise = hybridSearch(engine, opts.question, { limit: gatherLimit, expansion: false, // think provides its own anchor + graph context; no need for re-expansion + ...sourceScope, }).catch((e) => { process.stderr.write(`[think.gather] hybrid stream failed: ${(e as Error).message}\n`); return [] as SearchResult[]; @@ -119,6 +128,7 @@ export async function runGather( const takesKwPromise = engine.searchTakes(opts.question, { limit: takesLimit, takesHoldersAllowList: opts.takesHoldersAllowList, + ...sourceScope, }).catch((e) => { process.stderr.write(`[think.gather] takes-keyword stream failed: ${(e as Error).message}\n`); return [] as TakeHit[]; @@ -129,6 +139,7 @@ export async function runGather( ? engine.searchTakesVector(opts.questionEmbedding, { limit: takesLimit, takesHoldersAllowList: opts.takesHoldersAllowList, + ...sourceScope, }).catch((e) => { process.stderr.write(`[think.gather] takes-vector stream failed: ${(e as Error).message}\n`); return [] as TakeHit[]; @@ -137,7 +148,7 @@ export async function runGather( // Stream 4: graph walk (anchor only). const graphPromise: Promise = opts.anchor - ? engine.traversePaths(opts.anchor, { depth: graphDepth, direction: 'both' }) + ? engine.traversePaths(opts.anchor, { depth: graphDepth, direction: 'both', ...sourceScope }) .then(paths => { const slugs = new Set([opts.anchor!]); for (const p of paths) { diff --git a/src/core/think/index.ts b/src/core/think/index.ts index 4508631b1..8f3ab94cc 100644 --- a/src/core/think/index.ts +++ b/src/core/think/index.ts @@ -270,6 +270,8 @@ export async function runThink( anchor: opts.anchor, questionEmbedding, takesHoldersAllowList: opts.takesHoldersAllowList, + ...(opts.sourceId !== undefined ? { sourceId: opts.sourceId } : {}), + ...(opts.allowedSources !== undefined ? { sourceIds: opts.allowedSources } : {}), }); // Render evidence blocks for the prompt diff --git a/test/e2e/think-source-isolation-pglite.test.ts b/test/e2e/think-source-isolation-pglite.test.ts new file mode 100644 index 000000000..8d64ff4ac --- /dev/null +++ b/test/e2e/think-source-isolation-pglite.test.ts @@ -0,0 +1,92 @@ +import { afterAll, beforeAll, describe, expect, test } from 'bun:test'; +import { PGLiteEngine } from '../../src/core/pglite-engine.ts'; +import { runGather } from '../../src/core/think/gather.ts'; + +let engine: PGLiteEngine; + +beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); + + for (const sourceId of ['think-a', 'think-b', 'think-denied']) { + await engine.executeRaw( + `INSERT INTO sources (id, name, config) VALUES ($1, $1, '{}'::jsonb) ON CONFLICT DO NOTHING`, + [sourceId], + ); + } + + const fixtures = [ + ['think-a', 'people/think-anchor', 'authorized thinkscope anchor'], + ['think-b', 'people/think-allowed', 'authorized thinkscope evidence'], + ['think-denied', 'people/think-denied', 'denied thinkscope evidence'], + ] as const; + const takeVector = new Float32Array(1536).fill(0.01); + for (const [sourceId, slug, body] of fixtures) { + const page = await engine.putPage(slug, { + type: 'person', title: slug, compiled_truth: body, timeline: '', frontmatter: {}, + }, { sourceId }); + await engine.upsertChunks(slug, [{ + chunk_index: 0, chunk_text: body, chunk_source: 'compiled_truth', token_count: 4, + }], { sourceId }); + await engine.addTakesBatch([{ + page_id: page.id, row_num: 1, claim: 'thinkscope evidence', + kind: 'fact', holder: 'world', weight: 1, + }]); + await engine.executeRaw( + `UPDATE takes SET embedding = $1::vector WHERE page_id = $2`, + [`[${Array.from(takeVector).join(',')}]`, page.id], + ); + } + + await engine.addLink( + 'people/think-anchor', 'people/think-allowed', '', 'related', 'manual', + undefined, undefined, { fromSourceId: 'think-a', toSourceId: 'think-b' }, + ); + await engine.addLink( + 'people/think-anchor', 'people/think-denied', '', 'related', 'manual', + undefined, undefined, { fromSourceId: 'think-a', toSourceId: 'think-denied' }, + ); +}); + +afterAll(async () => { + await engine.disconnect(); +}); + +describe('think gather source isolation (#2200)', () => { + test('federated scope reaches hybrid, takes keyword/vector, and graph traversal', async () => { + const result = await runGather(engine, { + question: 'thinkscope evidence', + anchor: 'people/think-anchor', + questionEmbedding: new Float32Array(1536).fill(0.01), + sourceIds: ['think-a', 'think-b'], + gatherLimit: 50, + takesLimit: 50, + graphDepth: 2, + }); + + expect(result.pages.some(row => row.source_id === 'think-b')).toBe(true); + expect(result.pages.every(row => row.source_id !== 'think-denied')).toBe(true); + expect(result.takes.some(row => row.page_slug === 'people/think-allowed')).toBe(true); + expect(result.takes.every(row => row.page_slug !== 'people/think-denied')).toBe(true); + expect(result.graphSlugs).toContain('people/think-allowed'); + expect(result.graphSlugs).not.toContain('people/think-denied'); + }, 20_000); + + test('scalar sourceId reaches every gather stream', async () => { + const result = await runGather(engine, { + question: 'thinkscope evidence', + anchor: 'people/think-anchor', + questionEmbedding: new Float32Array(1536).fill(0.01), + sourceId: 'think-a', + gatherLimit: 50, + takesLimit: 50, + graphDepth: 2, + }); + + expect(result.pages.every(row => row.source_id === 'think-a')).toBe(true); + expect(result.takes.every(row => row.page_slug === 'people/think-anchor')).toBe(true); + expect(result.graphSlugs).not.toContain('people/think-allowed'); + expect(result.graphSlugs).not.toContain('people/think-denied'); + }, 20_000); +}); diff --git a/test/source-scope-resolver.test.ts b/test/source-scope-resolver.test.ts index e5b9d3671..9ee33d21c 100644 --- a/test/source-scope-resolver.test.ts +++ b/test/source-scope-resolver.test.ts @@ -12,6 +12,7 @@ import { describe, test, expect } from 'bun:test'; import { resolveRequestedScope, resolveCodeIntelScope, + thinkSourceScopeOpts, OperationError, type OperationContext, } from '../src/core/operations.ts'; @@ -56,6 +57,20 @@ describe('resolveRequestedScope — __all__ / all_sources', () => { }); }); +describe('think operation → runThink scope propagation', () => { + test('maps scalar sourceId without widening', () => { + expect(thinkSourceScopeOpts(ctxOf({ sourceId: 'tenant-a' }))).toEqual({ sourceId: 'tenant-a' }); + }); + + test('maps federated sourceIds to runThink allowedSources and wins over scalar', () => { + const ctx = ctxOf({ + sourceId: 'tenant-a', + auth: { allowedSources: ['tenant-a', 'tenant-b'] } as OperationContext['auth'], + }); + expect(thinkSourceScopeOpts(ctx)).toEqual({ allowedSources: ['tenant-a', 'tenant-b'] }); + }); +}); + describe('resolveRequestedScope — explicit source_id', () => { test('remote + explicit source_id OUTSIDE the grant is rejected', () => { const ctx = ctxOf({ remote: true, auth: { token: 't', clientId: 'c', scopes: [], allowedSources: ['a'] } as any });