Skip to content

Commit 617ce71

Browse files
committed
fix(knowledge): preserve fallback ranking and measure result bytes
1 parent cd1b6d8 commit 617ce71

8 files changed

Lines changed: 240 additions & 25 deletions

File tree

‎apps/sim/lib/copilot/request/tools/executor.test.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,42 @@ function buildPendingToolCall(): ToolCallState {
146146
}
147147
}
148148

149+
describe('tool result size diagnostics', () => {
150+
beforeEach(() => {
151+
vi.clearAllMocks()
152+
completeAsyncToolCall.mockResolvedValue(null)
153+
markAsyncToolRunning.mockResolvedValue(null)
154+
upsertAsyncToolCall.mockResolvedValue(null)
155+
})
156+
157+
it.each(['é🔎', { content: 'é🔎' }])(
158+
'records UTF-8 bytes after result projection for %j',
159+
async (output) => {
160+
executeTool.mockResolvedValueOnce({ success: true, output })
161+
const toolCall = buildPendingToolCall()
162+
const context = buildStreamingContext(toolCall)
163+
const endSpan = vi.spyOn(context.trace, 'endSpan')
164+
165+
const completion = await executeToolAndReport(toolCall.id, context, {
166+
userId: 'user-1',
167+
workflowId: 'workflow-1',
168+
resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
169+
})
170+
171+
expect(completion.status).toBe(MothershipStreamV1ToolOutcome.success)
172+
const serialized =
173+
typeof completion.data === 'string' ? completion.data : JSON.stringify(completion.data)
174+
expect(endSpan).toHaveBeenCalledWith(
175+
expect.objectContaining({
176+
kind: 'tool.execute',
177+
attributes: expect.objectContaining({ outputBytes: Buffer.byteLength(serialized) }),
178+
}),
179+
'ok'
180+
)
181+
}
182+
)
183+
})
184+
149185
describe('toolWatchdogTimeoutMs', () => {
150186
it('gives request-scoped MCP tools the long-running watchdog', () => {
151187
expect(toolWatchdogTimeoutMs('mcp-363de040-web_search_exa')).toBe(TOOL_WATCHDOG_LONG_RUNNING_MS)

‎apps/sim/lib/copilot/request/tools/executor.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -125,11 +125,11 @@ function summarizeToolResultForSpan(result: {
125125
const output = (result as { output: unknown }).output
126126
if (typeof output === 'string') {
127127
summary.outputKind = 'string'
128-
summary.outputBytes = output.length
128+
summary.outputBytes = Buffer.byteLength(output)
129129
} else if (output && typeof output === 'object') {
130130
summary.outputKind = Array.isArray(output) ? 'array' : 'object'
131131
try {
132-
summary.outputBytes = JSON.stringify(output).length
132+
summary.outputBytes = Buffer.byteLength(JSON.stringify(output))
133133
} catch {
134134
summary.outputBytes = 0
135135
}
@@ -143,7 +143,7 @@ function summarizeToolResultForSpan(result: {
143143
}
144144
} else if (output !== undefined && output !== null) {
145145
summary.outputKind = typeof output
146-
summary.outputBytes = String(output).length
146+
summary.outputBytes = Buffer.byteLength(String(output))
147147
}
148148
return summary
149149
}

‎apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11
/** @vitest-environment node */
22
import { beforeEach, describe, expect, it, vi } from 'vitest'
33

4-
const mocks = vi.hoisted(() => ({ search: vi.fn(), read: vi.fn(), authorizeChat: vi.fn() }))
4+
const mocks = vi.hoisted(() => ({
5+
search: vi.fn(),
6+
read: vi.fn(),
7+
authorizeChat: vi.fn(),
8+
info: vi.fn(),
9+
}))
10+
vi.mock('@sim/logger', () => ({
11+
createLogger: () => ({ info: mocks.info, error: vi.fn(), warn: vi.fn() }),
12+
}))
513
vi.mock('@/lib/copilot/chat/organization-chats', () => ({
614
authorizeOrganizationChatDelegation: { execute: mocks.authorizeChat },
715
}))
@@ -179,6 +187,48 @@ describe('Assistant retrieval tools', () => {
179187
})
180188
)
181189
})
190+
it.each([0, 20, 50])(
191+
'measures UTF-8 bytes for %i passages without logging their content',
192+
async (count) => {
193+
const content = 'Confidential passage é🔎'.repeat(100)
194+
mocks.search.mockResolvedValueOnce({
195+
knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
196+
results: Array.from({ length: count }, (_, index) => ({
197+
knowledgeBaseId: 'index',
198+
documentId: `doc-${index % 4}`,
199+
documentName: 'Private title',
200+
sourceUrl: null,
201+
sourceModifiedAt: null,
202+
metadata: {},
203+
content,
204+
chunkIndex: index,
205+
similarity: 1,
206+
})),
207+
})
208+
209+
const output = await searchWorkspaceServerTool.execute(
210+
{ query: 'Private query', ...(count === 50 ? { topK: 50 } : {}) },
211+
context
212+
)
213+
214+
expect(output.success).toBe(true)
215+
expect(mocks.info).toHaveBeenCalledWith(
216+
'Knowledge search completed',
217+
expect.objectContaining({
218+
toolCallId: 'call',
219+
toolResultBytes: Buffer.byteLength(JSON.stringify(output)),
220+
passageBytes: count * Buffer.byteLength(content),
221+
maxPassageBytes: count ? Buffer.byteLength(content) : 0,
222+
uniqueDocumentCount: Math.min(count, 4),
223+
})
224+
)
225+
const logged = JSON.stringify(mocks.info.mock.calls)
226+
expect(logged).not.toContain('Confidential passage')
227+
expect(logged).not.toContain('Private title')
228+
expect(logged).not.toContain('Private query')
229+
}
230+
)
231+
182232
it('returns stable citation IDs with internal links for uploaded documents', async () => {
183233
const result = await searchWorkspaceServerTool.execute({ query: 'orion' }, context)
184234
expect(result).toMatchObject({

‎apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
import { sourceAuthor } from '@/lib/knowledge/search/author'
1818
import { createKnowledgeDocumentCitation } from '@/lib/knowledge/search/citation'
1919
import {
20+
annotateSearchDiagnostics,
2021
measureSearchStage,
2122
recordSearchStageDuration,
2223
withSearchDiagnostics,
@@ -84,7 +85,7 @@ export const searchWorkspaceServerTool: BaseServerTool = {
8485
)
8586
return await measureSearchStage('tool_presentation', () => {
8687
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
87-
return {
88+
const output = {
8889
success: true,
8990
message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`,
9091
data: {
@@ -114,6 +115,14 @@ export const searchWorkspaceServerTool: BaseServerTool = {
114115
})),
115116
},
116117
}
118+
const passageBytes = output.data.results.map((item) => Buffer.byteLength(item.content))
119+
annotateSearchDiagnostics({
120+
toolResultBytes: Buffer.byteLength(JSON.stringify(output)),
121+
passageBytes: passageBytes.reduce((total, bytes) => total + bytes, 0),
122+
maxPassageBytes: Math.max(0, ...passageBytes),
123+
uniqueDocumentCount: new Set(output.data.results.map((item) => item.documentId)).size,
124+
})
125+
return output
117126
})
118127
} catch (error) {
119128
logger.error('Workspace search failed', { error })

‎apps/sim/lib/knowledge/__integration__/search-latency.integration.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,10 @@ const diagnosticSchema = z
132132
surface: z.enum(['dashboard', 'copilot']),
133133
outcome: z.literal('success'),
134134
elapsedMs: z.number(),
135+
toolResultBytes: z.number().int().nonnegative().optional(),
136+
passageBytes: z.number().int().nonnegative().optional(),
137+
maxPassageBytes: z.number().int().nonnegative().optional(),
138+
uniqueDocumentCount: z.number().int().nonnegative().optional(),
135139
stages: z.record(
136140
z.string(),
137141
z.object({
@@ -197,6 +201,15 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
197201
const diagnostics = diagnosticSchema.parse(completed[0][1])
198202
expect(diagnostics.stages.embedding.count).toBe(1)
199203
expect(diagnostics.stages.retrieval.count).toBe(1)
204+
if (diagnostics.surface === 'copilot') {
205+
const passageBytes = result.data.results.map((row) => Buffer.byteLength(row.content))
206+
expect(diagnostics.passageBytes).toBe(passageBytes.reduce((total, bytes) => total + bytes, 0))
207+
expect(diagnostics.maxPassageBytes).toBe(Math.max(0, ...passageBytes))
208+
expect(diagnostics.uniqueDocumentCount).toBe(
209+
new Set(result.data.results.map((row) => row.documentId)).size
210+
)
211+
expect(diagnostics.toolResultBytes).toBeGreaterThan(diagnostics.passageBytes!)
212+
}
200213
expect(captured.length).toBeLessThan(300)
201214
const searches = captured.filter(
202215
(item) =>

‎apps/sim/lib/knowledge/search/diagnostics.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,11 @@ export interface SearchDiagnosticMetadata {
6565
boostRecency?: boolean
6666
embeddingDimensions?: number
6767
resultCount?: number
68+
/** Tool output before the executor's final egress projection; counts only, never content. */
69+
toolResultBytes?: number
70+
passageBytes?: number
71+
maxPassageBytes?: number
72+
uniqueDocumentCount?: number
6873
}
6974

7075
interface StageTiming {

‎apps/sim/lib/knowledge/search/queries.test.ts‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,6 +606,76 @@ describe('live repository authorization follows ranked candidates', () => {
606606
)
607607
})
608608

609+
it('restarts exact ranking at zero and advances past already considered ANN candidates', async () => {
610+
const probe = Array.from({ length: 200 }, (_, index) =>
611+
candidate(`probe-${index}`, 'allowed-source')
612+
)
613+
const approximate = Array.from({ length: 20 }, (_, index) =>
614+
candidate(`approximate-${index}`, 'allowed-source')
615+
)
616+
queueTableRows(schemaMock.embedding, probe)
617+
queueTableRows(schemaMock.embedding, approximate)
618+
queueTableRows(schemaMock.embedding, [])
619+
queueTableRows(schemaMock.embedding, probe)
620+
queueTableRows(schemaMock.embedding, [])
621+
queueTableRows(schemaMock.embedding, approximate)
622+
queueTableRows(schemaMock.embedding, probe)
623+
queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
624+
queueTableRows(schemaMock.embedding, [
625+
{ id: 'selected', content: 'Reachable after the exact restart', distance: 0.1 },
626+
])
627+
628+
const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined })
629+
630+
expect(rows.map((row) => row.id)).toEqual(['selected'])
631+
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0], [20]])
632+
expect(getForConnectors).toHaveBeenCalledTimes(2)
633+
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(3)
634+
})
635+
636+
it('keeps the nearest exact results when an earlier ANN page hydrated only a farther result', async () => {
637+
const probe = Array.from({ length: 200 }, (_, index) =>
638+
candidate(`probe-${index}`, 'allowed-source')
639+
)
640+
queueTableRows(schemaMock.embedding, probe)
641+
queueTableRows(schemaMock.embedding, [
642+
{ ...candidate('far', 'allowed-source'), distance: 0.7 },
643+
...Array.from({ length: 19 }, (_, index) => candidate(`hidden-${index}`, 'allowed-source')),
644+
])
645+
queueTableRows(schemaMock.embedding, [{ id: 'far', content: 'Far result', distance: 0.7 }])
646+
queueTableRows(schemaMock.embedding, probe)
647+
queueTableRows(schemaMock.embedding, [])
648+
queueTableRows(schemaMock.embedding, [
649+
candidate('near', 'allowed-source'),
650+
candidate('nearer', 'allowed-source'),
651+
candidate('far', 'allowed-source'),
652+
])
653+
queueTableRows(schemaMock.embedding, [
654+
{ id: 'near', content: 'Near result', distance: 0.2 },
655+
{ id: 'nearer', content: 'Nearest result', distance: 0.1 },
656+
])
657+
658+
const rows = await handleVectorOnlySearch({
659+
...params,
660+
topK: 2,
661+
structuredFilters: undefined,
662+
})
663+
664+
expect(rows.map((row) => row.id)).toEqual(['nearer', 'near'])
665+
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0]])
666+
expect(
667+
hasMockCondition(
668+
dbChainMockFns.where.mock.calls.at(-1)![0],
669+
(node) =>
670+
node.type === 'inArray' &&
671+
node.column === schemaMock.embedding.id &&
672+
Array.isArray(node.values) &&
673+
node.values.length === 2 &&
674+
!node.values.includes('far')
675+
)
676+
).toBe(true)
677+
})
678+
609679
it.each(['vector', 'tag-vector', 'tags', 'keyword'] as const)(
610680
'%s ranks identifiers before verification and loads content under the full predicate',
611681
async (mode) => {

0 commit comments

Comments
 (0)