@@ -44,6 +44,7 @@ import {
4444 VECTOR_PROBE_DOCUMENT_LIMIT ,
4545 visibleDocumentsQuery ,
4646} from '@/lib/knowledge/search/queries'
47+ import { forgetIndexedVectorSources } from '@/lib/knowledge/search/source-vector-indexes'
4748import type { StructuredFilter } from '@/lib/knowledge/types'
4849
4950/**
@@ -1235,6 +1236,7 @@ describe('permitted-document planner', () => {
12351236 rerankRows = [ ]
12361237 sourceExactRows = [ ]
12371238 indexedSourceRows = [ ]
1239+ forgetIndexedVectorSources ( )
12381240 dbChainMockFns . execute . mockImplementation ( async ( query ) => {
12391241 const statement = render ( query ) . sql
12401242 if ( statement . includes ( 'pg_index' ) ) return indexedSourceRows
@@ -1274,13 +1276,61 @@ describe('permitted-document planner', () => {
12741276 traversedRows = [ { id : 'a' } ]
12751277 rerankRows = [ hit ( 'a' , null ) ]
12761278 queueTableRows ( schemaMock . embedding , [ hit ( 'a' , null ) ] )
1277- await handleVectorOnlySearch ( { ...params , permitted : { kind : 'unbounded' } } )
1279+ await handleVectorOnlySearch ( { ...params , permitted : { kind : 'unbounded' , broad : false } } )
12781280 const sqls = statements ( ) . map ( ( query ) => query . sql )
12791281 expect ( sqls . some ( ( sql ) => sql . includes ( 'AS visible' ) ) ) . toBe ( true )
12801282 expect ( sqls . some ( isProbeStatement ) ) . toBe ( false )
12811283 expect ( sqls . some ( isExactRanking ) ) . toBe ( false )
12821284 } )
12831285
1286+ it ( 'walks the whole graph once for a caller whose reach is broad' , async ( ) => {
1287+ const eligibility = { workspace : [ ] , admin : [ 'other-src' ] , members : [ 'member-src' ] }
1288+ indexedSourceRows = [ { name : 'idx' , connectorId : 'member-src' } ]
1289+ traversedRows = [ { id : 'walked-hit' , distance : 0.2 } ]
1290+ rerankRows = [ hit ( 'walked-hit' , 'member-src' ) ]
1291+ queueTableRows ( schemaMock . embedding , rerankRows )
1292+ await handleVectorOnlySearch ( {
1293+ ...params ,
1294+ permitted : { kind : 'unbounded' , broad : true } ,
1295+ accessPlan : {
1296+ connectors : eligibility ,
1297+ observers : { confirmed : [ 'm-1' ] , observed : [ ] } ,
1298+ memberSources : [ 'member-src' ] ,
1299+ } ,
1300+ } )
1301+ /** One walk over every source, scoped to the bases alone — no source is singled out. */
1302+ const walks = statements ( ) . filter ( ( query ) => query . sql . includes ( 'AS visible' ) )
1303+ expect ( walks ) . toHaveLength ( 1 )
1304+ expect ( JSON . stringify ( walks [ 0 ] ) ) . not . toContain ( '"right":"member-src"' )
1305+ expect ( statements ( ) . some ( ( query ) => query . sql . includes ( 'WITH readable_documents' ) ) ) . toBe ( false )
1306+ } )
1307+
1308+ it ( 'walks an indexed source a bounded caller is a member of instead of ranking it exactly' , async ( ) => {
1309+ const eligibility = { workspace : [ ] , admin : [ 'small-src' ] , members : [ 'member-src' ] }
1310+ indexedSourceRows = [ { name : 'idx' , connectorId : 'member-src' } ]
1311+ sourceExactRows = [ { id : 'small-hit' , distance : 0.3 , saturated : false } ]
1312+ traversedRows = [ { id : 'walked-hit' , distance : 0.2 } ]
1313+ rerankRows = [ hit ( 'walked-hit' , 'member-src' ) , hit ( 'small-hit' , 'small-src' ) ]
1314+ queueTableRows ( schemaMock . embedding , rerankRows )
1315+ await handleVectorOnlySearch ( {
1316+ ...params ,
1317+ topK : 2 ,
1318+ permitted : bounded (
1319+ { id : 'doc-a' , connectorId : 'member-src' } ,
1320+ { id : 'doc-b' , connectorId : 'small-src' }
1321+ ) ,
1322+ accessPlan : {
1323+ connectors : eligibility ,
1324+ observers : { confirmed : [ 'm-1' ] , observed : [ ] } ,
1325+ memberSources : [ 'member-src' ] ,
1326+ } ,
1327+ } )
1328+ const walks = statements ( ) . filter ( ( query ) => query . sql . includes ( 'AS visible' ) )
1329+ expect ( walks ) . toHaveLength ( 1 )
1330+ expect ( JSON . stringify ( walks [ 0 ] ) ) . toContain ( '"right":"member-src"' )
1331+ expect ( statements ( ) . some ( ( q ) => isExactRanking ( q . sql ) ) ) . toBe ( false )
1332+ } )
1333+
12841334 it ( 'walks a source the caller is a member of and ranks every other source exactly' , async ( ) => {
12851335 const eligibility = {
12861336 workspace : [ ] ,
@@ -1297,7 +1347,7 @@ describe('permitted-document planner', () => {
12971347 queueTableRows ( schemaMock . embedding , rerankRows )
12981348 await handleVectorOnlySearch ( {
12991349 ...params ,
1300- permitted : { kind : 'unbounded' } ,
1350+ permitted : { kind : 'unbounded' , broad : false } ,
13011351 accessPlan : { connectors : eligibility , observers, memberSources } ,
13021352 } )
13031353 const walks = statements ( ) . filter ( ( query ) => query . sql . includes ( 'AS visible' ) )
@@ -1320,7 +1370,7 @@ describe('permitted-document planner', () => {
13201370 queueTableRows ( schemaMock . embedding , rerankRows )
13211371 await handleVectorOnlySearch ( {
13221372 ...params ,
1323- permitted : { kind : 'unbounded' } ,
1373+ permitted : { kind : 'unbounded' , broad : false } ,
13241374 accessPlan : {
13251375 connectors : eligibility ,
13261376 observers : { confirmed : [ ] , observed : [ ] } ,
@@ -1347,7 +1397,7 @@ describe('permitted-document planner', () => {
13471397 await handleVectorOnlySearch ( {
13481398 ...params ,
13491399 topK : 2 ,
1350- permitted : { kind : 'unbounded' } ,
1400+ permitted : { kind : 'unbounded' , broad : false } ,
13511401 accessPlan : {
13521402 connectors : eligibility ,
13531403 observers : { confirmed : [ 'm-1' ] , observed : [ ] } ,
@@ -1368,7 +1418,7 @@ describe('permitted-document planner', () => {
13681418 queueTableRows ( schemaMock . embedding , rerankRows )
13691419 await handleVectorOnlySearch ( {
13701420 ...params ,
1371- permitted : { kind : 'unbounded' } ,
1421+ permitted : { kind : 'unbounded' , broad : false } ,
13721422 accessPlan : {
13731423 connectors : { workspace : [ ] , admin : [ 'sliced-src' ] , members : [ ] } ,
13741424 observers : { confirmed : [ ] , observed : [ ] } ,
@@ -1547,6 +1597,29 @@ describe('permitted-document planner', () => {
15471597 resolvePermittedDocuments ( { knowledgeBaseIds, access } )
15481598 const probes = ( ) => statements ( ) . filter ( ( query ) => isProbeStatement ( query . sql ) ) . length
15491599
1600+ it ( 'counts a saturated reach against the broad bound once, and remembers the answer' , async ( ) => {
1601+ /** The index holds a million documents; the bound is a quarter of them. */
1602+ const counts = { index : 1_000_000 , reached : 250_000 }
1603+ dbChainMockFns . execute . mockImplementation ( async ( query ) => {
1604+ const statement = render ( query ) . sql
1605+ if ( isProbeStatement ( statement ) ) return [ { id : null , connectorId : null , saturated : true } ]
1606+ if ( statement . includes ( ') reached' ) ) return [ { n : counts . reached } ]
1607+ if ( statement . includes ( 'count(*) AS n' ) ) return [ { n : counts . index } ]
1608+ return [ ]
1609+ } )
1610+ const reachCounts = ( ) => statements ( ) . filter ( ( query ) => query . sql . includes ( ') reached' ) )
1611+ const broad = await resolve ( scope ( 'broad-reach' ) )
1612+ expect ( broad ) . toEqual ( { kind : 'unbounded' , broad : true } )
1613+ expect ( reachCounts ( ) ) . toHaveLength ( 1 )
1614+ expect ( JSON . stringify ( reachCounts ( ) [ 0 ] ) ) . toContain ( '250000' )
1615+ await resolve ( scope ( 'broad-reach' ) )
1616+ expect ( reachCounts ( ) ) . toHaveLength ( 1 )
1617+ counts . reached = 120_000
1618+ const narrow = await resolve ( scope ( 'narrow-reach' ) )
1619+ expect ( narrow ) . toEqual ( { kind : 'unbounded' , broad : false } )
1620+ expect ( reachCounts ( ) ) . toHaveLength ( 2 )
1621+ } )
1622+
15501623 it ( 'is remembered, so a broad caller skips the probe on the next search' , async ( ) => {
15511624 probeRows = [ { id : null , connectorId : null , saturated : true } ]
15521625 const broad = scope ( 'broad' )
0 commit comments