Skip to content

Commit c4f0815

Browse files
icecrasher321icecrasher321
authored
fix(knowledge): optimize search and trace pipeline latency (#7782)
* fix(knowledge): optimize search and trace pipeline latency * docs: remove standalone search latency notes * fix(knowledge): preserve fallback ranking and measure result bytes
1 parent f4a86ee commit c4f0815

18 files changed

Lines changed: 1782 additions & 291 deletions

‎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: 96 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,12 @@ import {
1616
} from '@/lib/knowledge/application/workspace-search'
1717
import { sourceAuthor } from '@/lib/knowledge/search/author'
1818
import { createKnowledgeDocumentCitation } from '@/lib/knowledge/search/citation'
19+
import {
20+
annotateSearchDiagnostics,
21+
measureSearchStage,
22+
recordSearchStageDuration,
23+
withSearchDiagnostics,
24+
} from '@/lib/knowledge/search/diagnostics'
1925
import { intersectWorkspaceSearchFilters } from '@/lib/knowledge/search/filters'
2026
import { connectorDisplayName } from '@/lib/sim-search/connectors'
2127
import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection'
@@ -37,77 +43,99 @@ const CITATION_INSTRUCTION =
3743
export const searchWorkspaceServerTool: BaseServerTool = {
3844
name: 'search_workspace',
3945
async execute(raw, context?: ServerToolContext) {
40-
try {
41-
const scope = requireCopilotKnowledgeScope(context)
42-
const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw)
43-
const registry = context?.resolvedSecretTraceRegistry
44-
if (!registry) throw new Error('Knowledge result provenance is unavailable')
45-
const projected = projectResolvedSecretModelContent(query, registry)
46-
if (!projected.safe || typeof projected.value !== 'string') {
47-
return {
48-
success: false,
49-
message: 'Search query contains protected content. Rephrase the query.',
50-
}
51-
}
52-
const input = {
53-
query: projected.value,
54-
topK,
55-
filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch),
46+
return withSearchDiagnostics(
47+
{
5648
surface: context?.searchSurface ?? 'copilot',
57-
resultSecretRegistry: registry,
58-
signal: context?.abortSignal,
59-
} as const
60-
const result =
61-
scope.kind === 'organization'
62-
? await executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, {
63-
...input,
64-
organizationId: scope.organizationId,
65-
})
66-
: await executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, {
67-
...input,
68-
workspaceId: scope.workspaceId,
49+
toolCallId: context?.toolCallId,
50+
executionId: context?.executionId,
51+
},
52+
async () => {
53+
try {
54+
const inputStarted = performance.now()
55+
const scope = requireCopilotKnowledgeScope(context)
56+
const { query, topK, ...requestedFilters } = searchInputSchema.parse(raw)
57+
const registry = context?.resolvedSecretTraceRegistry
58+
if (!registry) throw new Error('Knowledge result provenance is unavailable')
59+
const projected = projectResolvedSecretModelContent(query, registry)
60+
if (!projected.safe || typeof projected.value !== 'string') {
61+
return {
62+
success: false,
63+
message: 'Search query contains protected content. Rephrase the query.',
64+
}
65+
}
66+
const input = {
67+
query: projected.value,
68+
topK,
69+
filters: intersectWorkspaceSearchFilters(requestedFilters, context?.assistantSearch),
70+
surface: context?.searchSurface ?? 'copilot',
71+
resultSecretRegistry: registry,
72+
signal: context?.abortSignal,
73+
} as const
74+
recordSearchStageDuration('tool_input', performance.now() - inputStarted)
75+
const result = await measureSearchStage('tool_application', () =>
76+
scope.kind === 'organization'
77+
? executeCopilotOrganizationKnowledgeUseCase(context, searchOrganizationKnowledge, {
78+
...input,
79+
organizationId: scope.organizationId,
80+
})
81+
: executeCopilotKnowledgeUseCase(context, searchWorkspaceKnowledge, {
82+
...input,
83+
workspaceId: scope.workspaceId,
84+
})
85+
)
86+
return await measureSearchStage('tool_presentation', () => {
87+
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
88+
const output = {
89+
success: true,
90+
message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`,
91+
data: {
92+
query,
93+
results: result.results.map((item) => ({
94+
documentId: item.documentId,
95+
knowledgeBaseId: item.knowledgeBaseId,
96+
knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '',
97+
siteName: item.connectorType
98+
? connectorDisplayName(item.connectorType)
99+
: names.get(item.knowledgeBaseId),
100+
documentName: item.documentName,
101+
sourceUrl: item.sourceUrl,
102+
connectorType: item.connectorType,
103+
sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null,
104+
author: sourceAuthor(item.metadata),
105+
content: item.content,
106+
chunkIndex: item.chunkIndex,
107+
similarity: item.similarity,
108+
...createKnowledgeDocumentCitation({
109+
scope,
110+
knowledgeBaseId: item.knowledgeBaseId,
111+
documentId: item.documentId,
112+
sourceUrl: item.sourceUrl,
113+
baseUrl: getBaseUrl(),
114+
}),
115+
})),
116+
},
117+
}
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,
69124
})
70-
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
71-
return {
72-
success: true,
73-
message: `Found ${result.results.length} passages. ${CITATION_INSTRUCTION}`,
74-
data: {
75-
query,
76-
results: result.results.map((item) => ({
77-
documentId: item.documentId,
78-
knowledgeBaseId: item.knowledgeBaseId,
79-
knowledgeBaseName: names.get(item.knowledgeBaseId) ?? '',
80-
siteName: item.connectorType
81-
? connectorDisplayName(item.connectorType)
82-
: names.get(item.knowledgeBaseId),
83-
documentName: item.documentName,
84-
sourceUrl: item.sourceUrl,
85-
connectorType: item.connectorType,
86-
sourceModifiedAt: item.sourceModifiedAt?.toISOString() ?? null,
87-
author: sourceAuthor(item.metadata),
88-
content: item.content,
89-
chunkIndex: item.chunkIndex,
90-
similarity: item.similarity,
91-
...createKnowledgeDocumentCitation({
92-
scope,
93-
knowledgeBaseId: item.knowledgeBaseId,
94-
documentId: item.documentId,
95-
sourceUrl: item.sourceUrl,
96-
baseUrl: getBaseUrl(),
97-
}),
98-
})),
99-
},
100-
}
101-
} catch (error) {
102-
logger.error('Workspace search failed', { error })
103-
return {
104-
success: false,
105-
message:
106-
error instanceof z.ZodError
107-
? 'Invalid search arguments'
108-
: messageForCopilotKnowledgeError(error),
125+
return output
126+
})
127+
} catch (error) {
128+
logger.error('Workspace search failed', { error })
129+
return {
130+
success: false,
131+
message:
132+
error instanceof z.ZodError
133+
? 'Invalid search arguments'
134+
: messageForCopilotKnowledgeError(error),
135+
}
136+
}
109137
}
110-
}
138+
)
111139
},
112140
}
113141

‎apps/sim/lib/knowledge/__integration__/application-acl.integration.ts‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,10 @@ import {
6262
seedKnowledgeMemberFixture,
6363
} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
6464
import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions'
65-
import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate'
65+
import {
66+
knowledgeAccessCondition,
67+
knowledgeMetadataCandidateAccessCondition,
68+
} from '@/lib/knowledge/access/predicate'
6669
import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope'
6770
import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks'
6871
import { readKnowledgeDocument } from '@/lib/knowledge/application/documents'
@@ -299,9 +302,16 @@ describe('indexed source content through real application access', () => {
299302
return result.results.map((row) => row.documentId)
300303
}
301304

302-
it.each(['workspace', 'admin', 'members'] as const)(
303-
'allows a remaining workspace ACL only in workspace mode, not during a %s transition',
304-
async (accessMode) => {
305+
it.each([
306+
['workspace', true],
307+
['workspace', false],
308+
['admin', true],
309+
['admin', false],
310+
['members', true],
311+
['members', false],
312+
] as const)(
313+
'requires settled workspace mode for a remaining workspace ACL (%s, rewrite pending=%s)',
314+
async (accessMode, accessRewritePending) => {
305315
const [savedConnector] = await db
306316
.select({
307317
accessMode: knowledgeConnector.accessMode,
@@ -324,18 +334,25 @@ describe('indexed source content through real application access', () => {
324334
.where(eq(document.id, documentId))
325335
await db
326336
.update(knowledgeConnector)
327-
.set({ accessMode, accessRewritePending: true })
337+
.set({ accessMode, accessRewritePending })
328338
.where(eq(knowledgeConnector.id, connectorId))
329-
const visible = await db
330-
.select({ id: document.id })
331-
.from(document)
332-
.where(
333-
and(
334-
eq(document.id, documentId),
335-
knowledgeAccessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] })
339+
for (const accessCondition of [
340+
knowledgeMetadataCandidateAccessCondition,
341+
knowledgeAccessCondition,
342+
]) {
343+
const visible = await db
344+
.select({ id: document.id })
345+
.from(document)
346+
.where(
347+
and(
348+
eq(document.id, documentId),
349+
accessCondition({ kind: 'workspace', tokens: ['pub', 'ws'] })
350+
)
336351
)
352+
expect(visible.map((row) => row.id)).toEqual(
353+
accessMode === 'workspace' && !accessRewritePending ? [documentId] : []
337354
)
338-
expect(visible.map((row) => row.id)).toEqual(accessMode === 'workspace' ? [documentId] : [])
355+
}
339356
} finally {
340357
await db
341358
.update(knowledgeConnector)

0 commit comments

Comments
 (0)