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

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

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<TakeHit[]> {
     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<TakeHit[]> {
     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[];
   }

src/core/postgres-engine.ts

@@ -4615,6 +4615,11 @@ export class PostgresEngine implements BrainEngine {
   async searchTakes(query: string, opts: SearchOpts & { takesHoldersAllowList?: string[] } = {}): Promise<TakeHit[]> {
     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}
     `;

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<string[]> = 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<string>([opts.anchor!]);
           for (const p of paths) {

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

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);
+});

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