diff --git a/apps/sim/lib/knowledge/mcp/read-budget.test.ts b/apps/sim/lib/knowledge/mcp/read-budget.test.ts new file mode 100644 index 00000000000..5d169390f83 --- /dev/null +++ b/apps/sim/lib/knowledge/mcp/read-budget.test.ts @@ -0,0 +1,73 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { shareReadBudget } from '@/lib/knowledge/mcp/read-budget' + +const read = (characters: number, text = 'x') => ({ + documentId: `doc-${characters}`, + chunks: [ + { + chunkIndex: 0, + content: text.repeat(characters), + startOffset: 0, + endOffset: characters, + totalCharacters: characters * 2, + }, + ], + hasMore: false, + next: null, +}) + +const lengths = (reads: { chunks: { content: string }[] }[]) => + reads.map((entry) => entry.chunks[0]?.content.length) + +describe('shareReadBudget', () => { + it('returns every read whole when the call fits its budget', () => { + const reads = [read(20_000), read(15_000), read(4_000)] + expect(shareReadBudget(reads)).toEqual(reads) + }) + + it('gives the budget short documents leave unused to the long one', () => { + const shared = shareReadBudget([read(20_000), read(20_000), ...Array(4).fill(read(1_000))]) + expect(lengths(shared)).toEqual([18_000, 18_000, 1_000, 1_000, 1_000, 1_000]) + }) + + it('cuts long documents to an equal share and continues each where it stops', () => { + const [first] = shareReadBudget([read(20_000), read(20_000), read(20_000), read(20_000)]) + expect(first?.chunks[0]).toMatchObject({ endOffset: 10_000 }) + expect(first?.chunks[0]?.content).toHaveLength(10_000) + expect(first).toMatchObject({ + hasMore: true, + next: { startChunkIndex: 0, startOffset: 10_000 }, + }) + }) + + it('never cuts a document below the standard page', () => { + expect(lengths(shareReadBudget(Array(6).fill(read(20_000))))).toEqual(Array(6).fill(8_000)) + }) + + it('counts text as JSON encodes it', () => { + const [quoted] = shareReadBudget([read(30_000, '"'), read(1_000), read(1_000)]) + const content = quoted?.chunks[0]?.content ?? '' + expect(JSON.stringify(content).length - 2).toBeLessThanOrEqual(38_000) + expect(content.length).toBeGreaterThan(18_000) + }) + + it('keeps the call within budget when a standard page encodes beyond its share', () => { + const shared = shareReadBudget([read(10_000, '"'), read(20_000), read(20_000), read(20_000)]) + const encoded = shared.reduce( + (total, entry) => total + JSON.stringify(entry.chunks[0]?.content ?? '').length - 2, + 0 + ) + expect(encoded).toBeLessThanOrEqual(40_000) + expect(lengths(shared)).toEqual([8_000, 8_000, 8_000, 8_000]) + }) + + it('passes a failed read through and leaves its share to the others', () => { + const failed = { documentId: 'missing', error: 'Not found' } + const shared = shareReadBudget([failed, read(20_000), read(20_000), read(20_000)]) + expect(shared[0]).toBe(failed) + expect(lengths(shared.slice(1) as ReturnType[])).toEqual([13_333, 13_333, 13_333]) + }) +}) diff --git a/apps/sim/lib/knowledge/mcp/read-budget.ts b/apps/sim/lib/knowledge/mcp/read-budget.ts new file mode 100644 index 00000000000..321de795380 --- /dev/null +++ b/apps/sim/lib/knowledge/mcp/read-budget.ts @@ -0,0 +1,79 @@ +import { encodedWindowEnd, READ_WINDOW_CHARACTERS } from '@/lib/sim-search/live/read-window' + +/** + * Characters of document text one batched read call returns. With about a thousand characters of + * metadata per document, it keeps the call 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. A call whose standard + * pages alone exceed it returns those pages, as a read without sharing would. + */ +const READ_CALL_CHARACTERS = 40_000 + +interface WindowedRead { + chunks: { content: string; startOffset: number; endOffset: number }[] + hasMore: boolean + next: { startChunkIndex: number; startOffset: number } | null +} + +const encodedLength = (text: string) => JSON.stringify(text).length - 2 + +const isWindowedRead = (document: object): document is WindowedRead => 'chunks' in document + +/** + * A read's text as JSON encodes it, and the part of that it always keeps: its standard page, or + * all of it when it is not a single chunk this budget can shorten. + */ +function encodedSize(read: object): { length: number; kept: number } { + if (!isWindowedRead(read)) return { length: 0, kept: 0 } + const length = read.chunks.reduce((total, chunk) => total + encodedLength(chunk.content), 0) + const [chunk] = read.chunks + if (!chunk || read.chunks.length !== 1) return { length, kept: length } + return { length, kept: encodedLength(chunk.content.slice(0, READ_WINDOW_CHARACTERS)) } +} + +/** + * Shares one call's text budget across the documents it read; an entry without text, such as a + * failed read, passes through. Every document keeps its standard page, and the budget those pages + * leave is shared as further text: documents needing less than an even share keep all theirs and + * leave the rest to the longer ones, which are cut to an equal share and continue through `next`. + */ +export function shareReadBudget(reads: readonly T[]): T[] { + const sizes = reads.map(encodedSize) + if (sizes.reduce((total, size) => total + size.length, 0) <= READ_CALL_CHARACTERS) + return [...reads] + const extras = sizes.map((size) => size.length - size.kept) + let remaining = READ_CALL_CHARACTERS - sizes.reduce((total, size) => total + size.kept, 0) + let left = reads.length + let cap = 0 + for (const extra of [...extras].sort((a, b) => a - b)) { + const share = Math.max(0, Math.floor(remaining / left)) + if (extra > share) { + cap = share + break + } + remaining -= extra + left-- + } + return reads.map((read, index) => + isWindowedRead(read) && extras[index] > cap ? cut(read, sizes[index].kept + cap) : read + ) +} + +/** A single-chunk read shortened so its text encodes within `budget`, continuing where it stops. */ +function cut(read: T, budget: number): T { + const [chunk] = read.chunks + if (!chunk || read.chunks.length !== 1) return read + const end = encodedWindowEnd( + chunk.content, + 0, + budget, + Math.min(chunk.content.length, READ_WINDOW_CHARACTERS) + ) + if (end >= chunk.content.length) return read + const endOffset = chunk.startOffset + end + return { + ...read, + chunks: [{ ...chunk, content: chunk.content.slice(0, end), endOffset }], + hasMore: true, + next: { startChunkIndex: 0, startOffset: endOffset }, + } +} diff --git a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts index e035aa87035..2e1dee80297 100644 --- a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts +++ b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts @@ -100,7 +100,7 @@ describe('Search MCP protocol', () => { expect(read.isError).not.toBe(true) expect(mocks.liveRead.mock.lastCall?.[0].input).toMatchObject({ limit: 8, - scaledWindow: { documentsInCall: 1 }, + scaledWindow: true, }) expect(read.content).toEqual([ { type: 'text', text: expect.stringContaining('Read evidence') }, @@ -112,7 +112,7 @@ describe('Search MCP protocol', () => { expect(batchRead.isError).not.toBe(true) expect(mocks.liveRead.mock.lastCall?.[0].input).toMatchObject({ limit: 8, - scaledWindow: { documentsInCall: 2 }, + scaledWindow: true, }) const offsetless = await client.callTool({ name: 'search', diff --git a/apps/sim/lib/knowledge/mcp/server.ts b/apps/sim/lib/knowledge/mcp/server.ts index 69f9074c401..096e896134f 100644 --- a/apps/sim/lib/knowledge/mcp/server.ts +++ b/apps/sim/lib/knowledge/mcp/server.ts @@ -22,6 +22,7 @@ import { recordOrganizationSearchMcpActivity, type SearchMcpActivityInput, } from '@/lib/knowledge/mcp/activity' +import { shareReadBudget } from '@/lib/knowledge/mcp/read-budget' import { KNOWLEDGE_MCP_READ_ONLY, type KnowledgeMcpToolRunner, @@ -298,7 +299,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe ...options, documentId: id, organizationId, - scaledWindow: { documentsInCall: documentId ? 1 : batch.length }, + scaledWindow: true, resultSecretRegistry: registry, signal, }, @@ -327,7 +328,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe } } ) - return projectResult({ documents }, registry) + return projectResult({ documents: shareReadBudget(documents) }, registry) }, Math.max(1, batch.length) ) diff --git a/apps/sim/lib/sim-search/live/application.test.ts b/apps/sim/lib/sim-search/live/application.test.ts index 472445c739e..4c3403096a4 100644 --- a/apps/sim/lib/sim-search/live/application.test.ts +++ b/apps/sim/lib/sim-search/live/application.test.ts @@ -1214,62 +1214,29 @@ describe('authorized live retrieval', () => { } } - 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('returns a longer window to a scaling caller, whatever its chunk limit', async () => { + for (const limit of [1, 3, 8]) { + const { length, next } = await readWindow({ limit, scaledWindow: true }) + expect(length).toBe(20_000) + expect(next).toEqual({ startChunkIndex: 0, startOffset: 20_000 }) + } }) 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) + const { length, next } = await readWindow({ scaledWindow: true, startOffset: 20_000 }) + expect(length).toBe(10_000) 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('keeps a quote-heavy window within its size once JSON-encoded', async () => { + const { length, encoded, next } = await readWindow({ scaledWindow: true }, '"'.repeat(30_000)) + expect(encoded).toBeLessThanOrEqual(20_000) + expect(length).toBeGreaterThanOrEqual(8_000) + expect(next).toEqual({ startChunkIndex: 0, startOffset: length }) }) - 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('keeps the standard window for a caller that does not scale', async () => { + expect((await readWindow({ limit: 8 })).length).toBe(8_000) }) }) it('refuses a chunk index other than zero before reading the document from its provider', async () => { diff --git a/apps/sim/lib/sim-search/live/application.ts b/apps/sim/lib/sim-search/live/application.ts index d5cf80f5e14..ca6415593db 100644 --- a/apps/sim/lib/sim-search/live/application.ts +++ b/apps/sim/lib/sim-search/live/application.ts @@ -75,6 +75,7 @@ import { loadLiveSearchPolicies } from '@/lib/sim-search/live/policy-store' import { LIVE_SEARCH_PROVIDER_IDS } from '@/lib/sim-search/live/provider-catalog' import { liveSearchGuidance } from '@/lib/sim-search/live/providers' import { LiveReadError } from '@/lib/sim-search/live/read-error' +import { encodedWindowEnd, READ_WINDOW_CHARACTERS } from '@/lib/sim-search/live/read-window' import { encodeLiveReference, includesLiveDocument, @@ -287,44 +288,25 @@ 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 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. + * Characters a scaled read returns for one document. Two documents at this size, with about a + * thousand characters of metadata each, stay 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; a call of more documents + * shares that budget, as the MCP server's batched read sizes it. */ -const SCALED_READ_CALL_CHARACTERS = 40_000 +const SCALED_READ_DOCUMENT_CHARACTERS = 20_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. + * Where one read's window ends: the standard page, or for a caller that scales its window, up to + * {@link SCALED_READ_DOCUMENT_CHARACTERS} as the result encodes it. A live document is one chunk, + * so its window is sized by characters, not by the chunk limit. The window is never shorter than + * the standard page. */ -function readWindowEnd( - content: string, - start: number, - { limit, scaledWindow }: Pick -): number { +function readWindowEnd(content: string, start: number, scaledWindow: boolean | undefined): 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 + return scaledWindow + ? encodedWindowEnd(content, start, SCALED_READ_DOCUMENT_CHARACTERS, standardEnd) + : standardEnd } /** @@ -1017,11 +999,8 @@ 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 } + /** Lengthens the window past the standard page, as {@link readWindowEnd} sizes it. */ + scaledWindow?: boolean startChunkIndex?: number startOffset?: number resultSecretRegistry: ResolvedSecretTraceRegistry @@ -1167,7 +1146,7 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({ 'validation', 'This document changed. Read again from the beginning' ) - const end = readWindowEnd(content, start, input) + const end = readWindowEnd(content, start, input.scaledWindow) return { documentId: input.documentId, knowledgeBaseId: '', diff --git a/apps/sim/lib/sim-search/live/read-window.ts b/apps/sim/lib/sim-search/live/read-window.ts new file mode 100644 index 00000000000..335f0d3c4df --- /dev/null +++ b/apps/sim/lib/sim-search/live/read-window.ts @@ -0,0 +1,22 @@ +/** Characters one standard read returns, the page budget indexed document reads use. */ +export const READ_WINDOW_CHARACTERS = 8000 + +/** + * Where a window of `content` from `start` ends so its text, as JSON encodes it, fits `budget` + * characters. JSON escapes quotes, backslashes and line breaks, so a window is shortened until its + * encoded text fits. It never ends before `minimumEnd`. + */ +export function encodedWindowEnd( + content: string, + start: number, + budget: number, + minimumEnd: number +): number { + let end = Math.min(content.length, start + budget) + for (let attempt = 0; attempt < 3 && end > minimumEnd; 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 minimumEnd +} diff --git a/apps/sim/lib/sim-search/live/term-fallback.test.ts b/apps/sim/lib/sim-search/live/term-fallback.test.ts index 34c1536789f..19bef424778 100644 --- a/apps/sim/lib/sim-search/live/term-fallback.test.ts +++ b/apps/sim/lib/sim-search/live/term-fallback.test.ts @@ -14,7 +14,7 @@ import type { NativeDocument, NativePage, NativeSearchInput } from '@/lib/sim-se const input: NativeSearchInput = { query: 'contractor offboarding runbook', limit: 4, scopes: [] } /** The any-term search's own budget, which a slow provider must not outlast. */ -const BUDGET_MS = 4000 +const BUDGET_MS = 2000 /** A caller that never cancels. */ const live = new AbortController().signal diff --git a/apps/sim/lib/sim-search/live/term-fallback.ts b/apps/sim/lib/sim-search/live/term-fallback.ts index b688e937a0d..08cb69d985a 100644 --- a/apps/sim/lib/sim-search/live/term-fallback.ts +++ b/apps/sim/lib/sim-search/live/term-fallback.ts @@ -17,9 +17,11 @@ const ANY_TERM_RESULTS = 5 /** * Time the any-term search may take before the all-terms page is returned without it, so a slow - * OR query cannot cost the account the matches it already has. + * OR query cannot cost the account the matches it already has. Confluence and Jira answer an + * any-term query within about 1.5 s, but Jira sometimes stalls until cancelled, and the whole + * search waits on this budget. */ -const ANY_TERM_BUDGET_MS = 4000 +const ANY_TERM_BUDGET_MS = 2000 const logger = createLogger('SimSearchTermFallback')