diff --git a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts index 62c34aa3dfc..e035aa87035 100644 --- a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts +++ b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts @@ -93,11 +93,27 @@ describe('Search MCP protocol', () => { expect(JSON.parse(searchContent.text)).toMatchObject({ results: [{ documentId, queryIndexes: [0] }], }) - const read = await client.callTool({ name: 'read_document', arguments: { documentId } }) + const read = await client.callTool({ + name: 'read_document', + arguments: { documentId, limit: 8 }, + }) expect(read.isError).not.toBe(true) + expect(mocks.liveRead.mock.lastCall?.[0].input).toMatchObject({ + limit: 8, + scaledWindow: { documentsInCall: 1 }, + }) expect(read.content).toEqual([ { type: 'text', text: expect.stringContaining('Read evidence') }, ]) + const batchRead = await client.callTool({ + name: 'read_document', + arguments: { documentIds: [documentId, `${documentId}b`], limit: 8 }, + }) + expect(batchRead.isError).not.toBe(true) + expect(mocks.liveRead.mock.lastCall?.[0].input).toMatchObject({ + limit: 8, + scaledWindow: { documentsInCall: 2 }, + }) const offsetless = await client.callTool({ name: 'search', arguments: { diff --git a/apps/sim/lib/knowledge/mcp/server.ts b/apps/sim/lib/knowledge/mcp/server.ts index facc4b8a0d7..69f9074c401 100644 --- a/apps/sim/lib/knowledge/mcp/server.ts +++ b/apps/sim/lib/knowledge/mcp/server.ts @@ -298,6 +298,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe ...options, documentId: id, organizationId, + scaledWindow: { documentsInCall: documentId ? 1 : batch.length }, resultSecretRegistry: registry, signal, }, diff --git a/apps/sim/lib/sim-search/live/application.test.ts b/apps/sim/lib/sim-search/live/application.test.ts index d2216b8e9ae..472445c739e 100644 --- a/apps/sim/lib/sim-search/live/application.test.ts +++ b/apps/sim/lib/sim-search/live/application.test.ts @@ -1189,6 +1189,89 @@ describe('authorized live retrieval', () => { }) ).rejects.toBe('user_stop:test') }) + describe('read window', () => { + const content = 'x'.repeat(30_000) + async function readWindow(extra: Record, text = content) { + const [result] = (await searchLiveKnowledge.execute({ principal, input })).results + if (!result) throw new Error('Expected the search fixture to return a result') + mocks.read.mockResolvedValue({ ...document, content: text }) + const read = await readLiveDocument.execute({ + principal, + input: { + workspaceId: 'workspace', + documentId: result.documentId, + limit: 3, + resultSecretRegistry: new ResolvedSecretTraceRegistry([]), + ...extra, + }, + }) + const [chunk] = read.chunks + if (!chunk) throw new Error('Expected the read to return a window') + return { + length: chunk.content.length, + encoded: JSON.stringify(chunk.content).length - 2, + next: read.next, + } + } + + const scaled = (documentsInCall: number) => ({ scaledWindow: { documentsInCall } }) + + it('returns a proportionally longer window when a scaling caller asks for more chunks', async () => { + const { length, next } = await readWindow({ limit: 8, ...scaled(1) }) + expect(length).toBe(21_333) + expect(next).toEqual({ startChunkIndex: 0, startOffset: 21_333 }) + }) + + it('continues a longer window from where the previous one ended', async () => { + const { length, next } = await readWindow({ limit: 8, ...scaled(1), startOffset: 21_333 }) + expect(length).toBe(30_000 - 21_333) + expect(next).toBeNull() + }) + + it.each([ + [5, 1, 13_333], + [8, 2, 20_000], + [8, 3, 13_333], + [8, 4, 10_000], + ])( + 'shares one call budget: limit %i across %i documents reads %i characters each', + async (limit, documentsInCall, expected) => { + expect((await readWindow({ limit, ...scaled(documentsInCall) })).length).toBe(expected) + } + ) + + it.each([ + [1, 21_333], + [2, 20_000], + ])( + 'keeps a quote-heavy window within its share once JSON-encoded, across %i documents', + async (documentsInCall, share) => { + const quotes = '"'.repeat(30_000) + const { length, encoded, next } = await readWindow( + { limit: 8, ...scaled(documentsInCall) }, + quotes + ) + expect(encoded).toBeLessThanOrEqual(share) + expect(length).toBeGreaterThanOrEqual(8_000) + expect(next).toEqual({ startChunkIndex: 0, startOffset: length }) + } + ) + + it('keeps the standard window for quote-heavy text at the default limit', async () => { + const { length } = await readWindow({ limit: 3, ...scaled(1) }, '"'.repeat(30_000)) + expect(length).toBe(8_000) + }) + + it.each([ + ['at the default limit', { limit: 3, ...scaled(1) }], + ['below the default limit', { limit: 1, ...scaled(1) }], + ['when the call budget is spread over five documents', { limit: 8, ...scaled(5) }], + ['when the call budget is spread over ten documents', { limit: 8, ...scaled(10) }], + ['for a caller that does not scale', { limit: 8 }], + ])('keeps the standard window %s', async (_, extra) => { + expect((await readWindow(extra)).length).toBe(8_000) + }) + }) it('refuses a chunk index other than zero before reading the document from its provider', async () => { const [result] = (await searchLiveKnowledge.execute({ principal, input })).results if (!result) throw new Error('Expected the search fixture to return a result') diff --git a/apps/sim/lib/sim-search/live/application.ts b/apps/sim/lib/sim-search/live/application.ts index 93e5e4461d5..d5cf80f5e14 100644 --- a/apps/sim/lib/sim-search/live/application.ts +++ b/apps/sim/lib/sim-search/live/application.ts @@ -287,8 +287,45 @@ export const LIVE_NOT_STARTED_MESSAGE = export const LIVE_OVER_ACCOUNT_LIMIT_MESSAGE = `Not searched: one search covers at most ${MAX_ACCOUNTS} accounts. Search this source or account alone.` /** Previews center on the query's longest matching term, like indexed passages. */ const PREVIEW_CHARACTERS = 1800 -/** Characters one read returns, the same page budget indexed document reads use. */ +/** Characters one read returns at the default limit, the page budget indexed document reads use. */ const READ_WINDOW_CHARACTERS = 8000 +/** The read limit {@link READ_WINDOW_CHARACTERS} is sized for, which is the default limit. */ +const READ_WINDOW_LIMIT = 3 +/** + * Characters the documents of one scaled read call share. Through four documents it keeps the + * call's result, with about a thousand characters of metadata per document, under the + * 50,000-character size at which MCP clients such as Claude Code save a tool result to a file + * rather than show it. From five documents the standard page per document is the larger, so + * those calls read exactly what they read without scaling. + */ +const SCALED_READ_CALL_CHARACTERS = 40_000 + +/** + * Where one read's window ends. A live document is one chunk, so a caller that scales the window + * gets a longer one when it asks for more chunks than the default, bounded by an even share of the + * call's budget across the documents it reads. The share counts the window as the result encodes + * it, since JSON escapes quotes, backslashes and line breaks, so a longer window is shortened until + * its encoded text fits. The window is never shorter than the standard page, which is also what a + * read at or below the default limit, or from any other caller, returns. + */ +function readWindowEnd( + content: string, + start: number, + { limit, scaledWindow }: Pick +): number { + const standardEnd = Math.min(content.length, start + READ_WINDOW_CHARACTERS) + if (!scaledWindow || limit <= READ_WINDOW_LIMIT) return standardEnd + const scaled = Math.floor((READ_WINDOW_CHARACTERS * limit) / READ_WINDOW_LIMIT) + const share = Math.floor(SCALED_READ_CALL_CHARACTERS / Math.max(1, scaledWindow.documentsInCall)) + const budget = Math.min(scaled, share) + let end = Math.min(content.length, start + budget) + for (let attempt = 0; attempt < 3 && end > standardEnd; attempt++) { + const encoded = JSON.stringify(content.slice(start, end)).length - 2 + if (encoded <= budget) return end + end = start + Math.floor(((end - start) * budget) / encoded) + } + return standardEnd +} /** * Verifies candidates in rank order. A revoked grant fails the whole account; any other @@ -980,6 +1017,11 @@ export type LiveReadInput = ResourceOwner & { documentId: string filters?: WorkspaceSearchFilters limit: number + /** + * Lengthens the window for a limit above the default, sharing one call's budget across the + * documents it reads, as {@link readWindowEnd} sizes it. + */ + scaledWindow?: { documentsInCall: number } startChunkIndex?: number startOffset?: number resultSecretRegistry: ResolvedSecretTraceRegistry @@ -1125,7 +1167,7 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({ 'validation', 'This document changed. Read again from the beginning' ) - const end = Math.min(content.length, start + READ_WINDOW_CHARACTERS) + const end = readWindowEnd(content, start, input) return { documentId: input.documentId, knowledgeBaseId: '',