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
73 changes: 73 additions & 0 deletions apps/sim/lib/knowledge/mcp/read-budget.test.ts
Original file line number Diff line number Diff line change
@@ -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<typeof read>[])).toEqual([13_333, 13_333, 13_333])
})
})
79 changes: 79 additions & 0 deletions apps/sim/lib/knowledge/mcp/read-budget.ts
Original file line number Diff line number Diff line change
@@ -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<T extends object>(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<T extends WindowedRead>(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 },
}
}
4 changes: 2 additions & 2 deletions apps/sim/lib/knowledge/mcp/server.protocol.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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') },
Expand All @@ -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',
Expand Down
5 changes: 3 additions & 2 deletions apps/sim/lib/knowledge/mcp/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
},
Expand Down Expand Up @@ -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)
)
Expand Down
63 changes: 15 additions & 48 deletions apps/sim/lib/sim-search/live/application.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down
55 changes: 17 additions & 38 deletions apps/sim/lib/sim-search/live/application.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<LiveReadInput, 'limit' | 'scaledWindow'>
): 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
}

/**
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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: '',
Expand Down
22 changes: 22 additions & 0 deletions apps/sim/lib/sim-search/live/read-window.ts
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 1 addition & 1 deletion apps/sim/lib/sim-search/live/term-fallback.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading