fix(knowledge): count list documents once instead of per 400-base batch by waleedlatif1 · Pull Request #7759 · simstudioai/sim · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions apps/sim/lib/knowledge/access/scope.test.ts
4 changes: 4 additions & 0 deletions apps/sim/lib/knowledge/access/scope.ts
Original file line number Diff line number Diff line change
Expand Up @@ -397,6 +397,10 @@ function createAccessProvider(
async get() {
return (await identity()).access
},
async hasLiveSourceReaders() {
const { access, githubReaders, confluenceReaders } = await identity()
return access.kind === 'user' && (githubReaders.length > 0 || confluenceReaders.length > 0)
},
async getForConnectors(connectorIds, signal) {
const ids = boundedIds(connectorIds)
const cancellation =
Expand Down
6 changes: 6 additions & 0 deletions apps/sim/lib/knowledge/access/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,12 @@ export interface KnowledgeAccessProvider {
documentIds: readonly string[],
signal?: AbortSignal
): Promise<KnowledgeAccessScope>
/**
* Whether the reader holds any credential a live source (GitHub, Confluence) could
* authorize beyond the stored ACL. Without one, candidate discovery can only re-prove
* the ordinary predicate, so readers skip it. Absent means unknown: discover.
*/
hasLiveSourceReaders?(): Promise<boolean>
}

/** Two existing search legs each contribute at most 200 candidates to one authorization batch. */
Expand Down
15 changes: 15 additions & 0 deletions apps/sim/lib/knowledge/read-access.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,21 @@ describe('knowledgeReadAccessBatches', () => {
expect(gt).toHaveBeenCalledWith(document.connectorId, first.at(-1)!.connectorId)
})

it('yields only the ordinary predicate for a reader without live-source credentials', async () => {
const resolve = vi.fn(async () => identity)
const provider: KnowledgeAccessProvider = {
get: async () => identity,
getForConnectors: resolve,
getForDocuments: async () => identity,
hasLiveSourceReaders: async () => false,
}
const batches = []
for await (const predicate of knowledgeReadAccessBatches(provider, [])) batches.push(predicate)
expect(batches).toHaveLength(1)
expect(resolve).not.toHaveBeenCalled()
expect(dbChainMockFns.selectDistinct).not.toHaveBeenCalled()
})

it('does not enumerate sources after a satisfied ordinary existence probe', async () => {
const resolve = vi.fn(async () => identity)
const provider: KnowledgeAccessProvider = {
Expand Down
1 change: 1 addition & 0 deletions apps/sim/lib/knowledge/read-access.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ export async function* knowledgeReadAccessBatches(
const ordinary = knowledgeAccessCondition(scope)
yield ordinary
if (!provider || scope.kind !== 'user') return
if (provider.hasLiveSourceReaders && !(await provider.hasLiveSourceReaders())) return

let cursor: string | undefined
while (true) {
Expand Down
55 changes: 52 additions & 3 deletions apps/sim/lib/knowledge/service.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -551,12 +551,11 @@ describe('knowledge base counts with live source permissions', () => {
id: 'kb-1',
workspaceId: 'ws-1',
chunkingConfig: {},
docCount: 99,
tokenCount: 999,
docCount: 2,
tokenCount: 10,
createdAt: new Date('2026-01-01'),
},
])
queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 }])
queueTableRows(schemaMock.document, [{ connectorId: 'confluence-source' }])
queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 3, tokenCount: 20 }])
const result = await getWorkspaceKnowledgeBases('ws-1', 'active', { access, limit: 2 })
Expand Down Expand Up @@ -587,6 +586,56 @@ describe('knowledge base counts with live source permissions', () => {
).toBe(true)
})

it('counts an unpaged list in one joined query and one discovery pass bounded by the list filter', async () => {
const { access, getForConnectors } = reader()
const bases = Array.from({ length: 1000 }, (_, index) => ({
id: `kb-${index}`,
workspaceId: 'ws-1',
chunkingConfig: {},
docCount: 1,
tokenCount: 1,
createdAt: new Date('2026-01-01'),
}))
queueTableRows(schemaMock.knowledgeBase, bases)
queueTableRows(schemaMock.document, [{ connectorId: 'confluence-source' }])
queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-7', docCount: 3, tokenCount: 20 }])
const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { access })
expect(result.data).toHaveLength(1000)
expect(result.data[7]).toMatchObject({ docCount: 4, tokenCount: 21 })
expect(result.data[8]).toMatchObject({ docCount: 1, tokenCount: 1 })
expect(getForConnectors).toHaveBeenCalledOnce()
expect(dbChainMockFns.selectDistinct).toHaveBeenCalledOnce()
expect(dbChainMockFns.groupBy).toHaveBeenCalledTimes(2)
expect(
dbChainMockFns.where.mock.calls.some(([condition]) =>
hasMockCondition(
condition,
(node) => node.type === 'inArray' && node.column === schemaMock.knowledgeBase.id
)
)
).toBe(false)
})

it('never discovers live sources for a reader without live-source credentials', async () => {
const { access, getForConnectors } = reader()
access.hasLiveSourceReaders = async () => false
queueTableRows(schemaMock.knowledgeBase, [
{
id: 'kb-1',
workspaceId: 'ws-1',
chunkingConfig: {},
docCount: 2,
tokenCount: 10,
createdAt: new Date('2026-01-01'),
},
])
const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { access })
expect(result.data[0]).toMatchObject({ docCount: 2, tokenCount: 10 })
expect(getForConnectors).not.toHaveBeenCalled()
expect(dbChainMockFns.selectDistinct).not.toHaveBeenCalled()
expect(dbChainMockFns.groupBy).toHaveBeenCalledOnce()
})

it('does not retain stale totals when a live source no longer authorizes its documents', async () => {
const { access, getForConnectors } = reader()
queueTableRows(schemaMock.document, [])
Expand Down
124 changes: 77 additions & 47 deletions apps/sim/lib/knowledge/service.ts
Loading