Skip to content

Commit 122ea0b

Browse files
fix(search): inject personal integrations into chat prompts (#7751)
1 parent 1674e83 commit 122ea0b

15 files changed

Lines changed: 343 additions & 174 deletions

File tree

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { beforeEach, describe, expect, it, vi } from 'vitest'
5+
6+
const { authorizeChat, listIntegrations } = vi.hoisted(() => ({
7+
authorizeChat: vi.fn(),
8+
listIntegrations: vi.fn(),
9+
}))
10+
11+
vi.mock('@/lib/copilot/chat/organization-chats', () => ({
12+
authorizeOrganizationChatDelegation: { execute: authorizeChat },
13+
}))
14+
vi.mock('@/lib/knowledge/application/personal-search-integrations', () => ({
15+
listPersonalSearchIntegrations: { execute: listIntegrations },
16+
}))
17+
18+
import { loadCopilotSearchIntegrations } from '@/lib/copilot/application/load-search-integrations'
19+
import type { listPersonalSearchIntegrations } from '@/lib/knowledge/application/personal-search-integrations'
20+
21+
type InventoryPage = Awaited<ReturnType<typeof listPersonalSearchIntegrations.execute>>
22+
23+
const context = {
24+
userId: 'person-1',
25+
organizationId: 'org-1',
26+
chatId: 'private-chat-1',
27+
messageId: 'message-1',
28+
}
29+
const emptyPage: InventoryPage = {
30+
connections: [],
31+
available: [],
32+
completedCredentialId: null,
33+
nextCursor: null,
34+
}
35+
36+
describe('loadCopilotSearchIntegrations', () => {
37+
beforeEach(() => {
38+
vi.clearAllMocks()
39+
authorizeChat.mockResolvedValue(undefined)
40+
listIntegrations.mockResolvedValue(emptyPage)
41+
})
42+
43+
it('authorizes the private chat and reads only for the authenticated person and organization', async () => {
44+
expect(await loadCopilotSearchIntegrations(context)).toBe('{"connections":[],"available":[]}')
45+
const principal = authorizeChat.mock.calls[0][0].principal
46+
expect(principal).toMatchObject({
47+
kind: 'organization_delegated',
48+
serviceId: 'copilot',
49+
subjectUserId: 'person-1',
50+
organizationId: 'org-1',
51+
delegationId: 'message-1',
52+
audience: 'sim:knowledge',
53+
resourceScope: { chatId: 'private-chat-1' },
54+
})
55+
expect(listIntegrations).toHaveBeenCalledExactlyOnceWith({
56+
principal,
57+
input: { organizationId: 'org-1' },
58+
})
59+
expect(authorizeChat.mock.invocationCallOrder[0]).toBeLessThan(
60+
listIntegrations.mock.invocationCallOrder[0]
61+
)
62+
})
63+
64+
it('loads every page and preserves account status and exact connection controls', async () => {
65+
const available: InventoryPage['available'][number] = {
66+
name: 'Gmail',
67+
description: 'Personal mail',
68+
target: { type: 'link', provider: 'google-email', connectorType: 'gmail' },
69+
}
70+
const connection: InventoryPage['connections'][number] = {
71+
name: 'Gmail',
72+
providerId: 'google-email',
73+
connectorType: 'gmail',
74+
connectorId: 'source-1',
75+
knowledgeBaseId: 'kb-1',
76+
description: 'Personal mail',
77+
accounts: [
78+
{
79+
credentialId: 'account-1',
80+
displayName: 'me@example.com',
81+
status: 'reconnect_needed',
82+
action: { ...available.target, connectorId: 'source-1', credentialId: 'account-1' },
83+
},
84+
],
85+
connectionStatus: 'reconnect_needed',
86+
indexingStatus: 'indexed',
87+
searchableDocuments: 7,
88+
action: null,
89+
}
90+
listIntegrations
91+
.mockResolvedValueOnce({ ...emptyPage, available: [available], nextCursor: 'page-2' })
92+
.mockResolvedValueOnce({ ...emptyPage, connections: [connection], available: [available] })
93+
94+
expect(JSON.parse(await loadCopilotSearchIntegrations(context))).toEqual({
95+
connections: [connection],
96+
available: [available],
97+
})
98+
expect(listIntegrations.mock.calls[1][0]).toEqual({
99+
principal: authorizeChat.mock.calls[0][0].principal,
100+
input: { organizationId: 'org-1', cursor: 'page-2' },
101+
})
102+
})
103+
104+
it('does not read inventory when private-chat authorization fails', async () => {
105+
authorizeChat.mockRejectedValueOnce(new Error('Chat belongs to another person'))
106+
await expect(loadCopilotSearchIntegrations(context)).rejects.toThrow(
107+
'Chat belongs to another person'
108+
)
109+
expect(listIntegrations).not.toHaveBeenCalled()
110+
})
111+
112+
it('fails the turn if a later page cannot be read', async () => {
113+
listIntegrations
114+
.mockResolvedValueOnce({ ...emptyPage, nextCursor: 'page-2' })
115+
.mockRejectedValueOnce(new Error('Inventory unavailable'))
116+
await expect(loadCopilotSearchIntegrations(context)).rejects.toThrow('Inventory unavailable')
117+
})
118+
119+
it('rejects a repeated cursor instead of looping or returning partial inventory', async () => {
120+
listIntegrations.mockResolvedValue({ ...emptyPage, nextCursor: 'page-2' })
121+
await expect(loadCopilotSearchIntegrations(context)).rejects.toThrow(
122+
'pagination did not advance'
123+
)
124+
expect(listIntegrations).toHaveBeenCalledTimes(2)
125+
})
126+
127+
it('bounds page loading when the inventory never ends', async () => {
128+
listIntegrations.mockImplementation(async () => ({
129+
...emptyPage,
130+
nextCursor: `page-${listIntegrations.mock.calls.length + 1}`,
131+
}))
132+
await expect(loadCopilotSearchIntegrations(context)).rejects.toThrow('pagination limit')
133+
expect(listIntegrations).toHaveBeenCalledTimes(100)
134+
})
135+
136+
it('rejects oversized prompt content instead of silently truncating it', async () => {
137+
listIntegrations.mockResolvedValueOnce({
138+
...emptyPage,
139+
available: [
140+
{
141+
name: 'Gmail',
142+
description: 'x'.repeat(256 * 1024),
143+
target: { type: 'link', provider: 'google-email', connectorType: 'gmail' },
144+
},
145+
],
146+
})
147+
await expect(loadCopilotSearchIntegrations(context)).rejects.toThrow('prompt size limit')
148+
})
149+
150+
it('stops loading when the turn is cancelled between pages', async () => {
151+
const controller = new AbortController()
152+
listIntegrations.mockImplementationOnce(async () => {
153+
controller.abort(new Error('Turn cancelled'))
154+
return { ...emptyPage, nextCursor: 'page-2' }
155+
})
156+
await expect(
157+
loadCopilotSearchIntegrations({ ...context, signal: controller.signal })
158+
).rejects.toThrow('Turn cancelled')
159+
expect(listIntegrations).toHaveBeenCalledTimes(1)
160+
})
161+
})
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
import {
2+
COPILOT_APPLICATION_DELEGATION_TTL_MS,
3+
createTrustedOrganizationCopilotPrincipal,
4+
} from '@/lib/copilot/auth/application-delegation'
5+
import { authorizeOrganizationChatDelegation } from '@/lib/copilot/chat/organization-chats'
6+
import { knowledgeDelegationPolicy } from '@/lib/knowledge/application/authorization'
7+
import { listPersonalSearchIntegrations } from '@/lib/knowledge/application/personal-search-integrations'
8+
9+
const MAX_INVENTORY_PAGES = 100
10+
const MAX_INVENTORY_BYTES = 256 * 1024
11+
12+
interface SearchIntegrationsContext {
13+
userId: string
14+
organizationId: string
15+
chatId: string
16+
messageId: string
17+
signal?: AbortSignal
18+
}
19+
20+
type IntegrationInventory = Awaited<ReturnType<typeof listPersonalSearchIntegrations.execute>>
21+
22+
/** Loads the complete current person's Search inventory for one authenticated chat turn. */
23+
export async function loadCopilotSearchIntegrations(
24+
context: SearchIntegrationsContext
25+
): Promise<string> {
26+
context.signal?.throwIfAborted()
27+
const principal = createTrustedOrganizationCopilotPrincipal(
28+
{ ...context, delegationId: context.messageId },
29+
{
30+
audience: knowledgeDelegationPolicy.audience,
31+
ttlMs: COPILOT_APPLICATION_DELEGATION_TTL_MS,
32+
}
33+
)
34+
await authorizeOrganizationChatDelegation.execute({ principal })
35+
36+
const connections: IntegrationInventory['connections'] = []
37+
const available = new Map<string, IntegrationInventory['available'][number]>()
38+
const cursors = new Set<string>()
39+
let cursor: string | undefined
40+
for (let pageNumber = 0; pageNumber < MAX_INVENTORY_PAGES; pageNumber++) {
41+
context.signal?.throwIfAborted()
42+
const page = await listPersonalSearchIntegrations.execute({
43+
principal,
44+
input: { organizationId: context.organizationId, ...(cursor ? { cursor } : {}) },
45+
})
46+
context.signal?.throwIfAborted()
47+
connections.push(...page.connections)
48+
for (const entry of page.available) {
49+
available.set(JSON.stringify(entry.target), entry)
50+
}
51+
const inventory = JSON.stringify({ connections, available: [...available.values()] })
52+
if (Buffer.byteLength(inventory) > MAX_INVENTORY_BYTES) {
53+
throw new Error('Search integration inventory exceeds the prompt size limit')
54+
}
55+
if (page.nextCursor === null) return inventory
56+
if (cursors.has(page.nextCursor)) {
57+
throw new Error('Search integration inventory pagination did not advance')
58+
}
59+
cursors.add(page.nextCursor)
60+
cursor = page.nextCursor
61+
}
62+
throw new Error('Search integration inventory exceeds the pagination limit')
63+
}

apps/sim/lib/copilot/assistant/tool-policy.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ import type { ToolMetadata } from '@/tools/metadata'
33
export const ASSISTANT_TOOLS = new Set([
44
'search_workspace',
55
'read_document',
6-
'list_integrations',
76
'search_integration_tools',
87
'call_integration_tool',
98
'oauth_get_auth_link',

apps/sim/lib/copilot/generated/tool-catalog-v1.ts

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@ export interface ToolCatalogEntry {
7474
| 'knowledge'
7575
| 'list_deployment_versions'
7676
| 'list_integration_tools'
77-
| 'list_integrations'
7877
| 'list_workspace_mcp_servers'
7978
| 'load_deployment'
8079
| 'load_integration_tool'
@@ -213,7 +212,6 @@ export interface ToolCatalogEntry {
213212
| 'knowledge'
214213
| 'list_deployment_versions'
215214
| 'list_integration_tools'
216-
| 'list_integrations'
217215
| 'list_workspace_mcp_servers'
218216
| 'load_deployment'
219217
| 'load_integration_tool'
@@ -3543,21 +3541,6 @@ export const ListIntegrationTools: ToolCatalogEntry = {
35433541
},
35443542
}
35453543

3546-
export const ListIntegrations: ToolCatalogEntry = {
3547-
id: 'list_integrations',
3548-
name: 'list_integrations',
3549-
route: 'sim',
3550-
mode: 'async',
3551-
parameters: {
3552-
additionalProperties: false,
3553-
properties: {
3554-
connectorType: { maxLength: 100, minLength: 1, type: 'string' },
3555-
cursor: { maxLength: 1024, minLength: 1, type: 'string' },
3556-
},
3557-
type: 'object',
3558-
},
3559-
}
3560-
35613544
export const ListWorkspaceMcpServers: ToolCatalogEntry = {
35623545
id: 'list_workspace_mcp_servers',
35633546
name: 'list_workspace_mcp_servers',
@@ -7664,7 +7647,6 @@ export const TOOL_CATALOG: Record<string, ToolCatalogEntry> = {
76647647
[Knowledge.id]: Knowledge,
76657648
[ListDeploymentVersions.id]: ListDeploymentVersions,
76667649
[ListIntegrationTools.id]: ListIntegrationTools,
7667-
[ListIntegrations.id]: ListIntegrations,
76687650
[ListWorkspaceMcpServers.id]: ListWorkspaceMcpServers,
76697651
[LoadDeployment.id]: LoadDeployment,
76707652
[LoadIntegrationTool.id]: LoadIntegrationTool,

apps/sim/lib/copilot/generated/tool-schemas-v1.ts

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3474,25 +3474,6 @@ export const TOOL_RUNTIME_SCHEMAS: Record<string, ToolRuntimeSchemaEntry> = {
34743474
},
34753475
resultSchema: undefined,
34763476
},
3477-
list_integrations: {
3478-
parameters: {
3479-
additionalProperties: false,
3480-
properties: {
3481-
connectorType: {
3482-
maxLength: 100,
3483-
minLength: 1,
3484-
type: 'string',
3485-
},
3486-
cursor: {
3487-
maxLength: 1024,
3488-
minLength: 1,
3489-
type: 'string',
3490-
},
3491-
},
3492-
type: 'object',
3493-
},
3494-
resultSchema: undefined,
3495-
},
34963477
list_workspace_mcp_servers: {
34973478
parameters: {
34983479
type: 'object',

0 commit comments

Comments
 (0)