Skip to content

Commit eb22479

Browse files
Sg312Sg312
committed
Fix GitHub connected-account tools in Search Assistant
1 parent 94a8d0a commit eb22479

10 files changed

Lines changed: 291 additions & 18 deletions

File tree

‎apps/sim/executor/utils/credential-token.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import { createLogger } from '@sim/logger'
22
import { AuthType } from '@/lib/auth/hybrid'
3+
import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags'
34
import { createCopilotManagedOAuthPrincipal } from '@/lib/credentials/application/copilot-managed-oauth-delegation'
45
import { bindExecutorManagedOAuthDelegation } from '@/lib/credentials/application/managed-oauth-delegation'
56
import { authorizePersonalCredential } from '@/lib/credentials/application/personal-credentials'
67
import { executeCopilotCredentialUseCase } from '@/lib/mothership/application/execute-credential-use-case'
78
import { resolveCopilotOrganizationPersonalToken } from '@/lib/mothership/application/resolve-organization-personal-token'
9+
import { projectAssistantConnectedAccountTool } from '@/lib/mothership/assistant/connected-account-tool'
810
import type { CopilotExecutionContext } from '@/lib/mothership/auth/application-delegation'
911
import {
1012
type CredentialTokenPayload,
@@ -61,7 +63,10 @@ export async function resolveExecutorCredentialToken(
6163
if (!userId || userId !== copilotExecutionContext.userId || executorDelegationOrigin) {
6264
throw new Error('Assistant credential use requires the authenticated person for this turn.')
6365
}
64-
const tool = toolId ? getToolMetadata(toolId) : undefined
66+
const original = toolId ? getToolMetadata(toolId) : undefined
67+
const tool = original
68+
? projectAssistantConnectedAccountTool(original, isLiveEnterpriseSearchEnabled)
69+
: undefined
6570
if (
6671
!tool?.oauth?.required ||
6772
(!copilotExecutionContext.workspaceId && !copilotExecutionContext.organizationId) ||

‎apps/sim/lib/credentials/application/resolve-organization-personal-token.test.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
/** @vitest-environment node */
2+
23
import type { OrganizationDelegatedPrincipal } from '@sim/auth/principal'
4+
import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
35
import { beforeEach, describe, expect, it, vi } from 'vitest'
46

57
const mocks = vi.hoisted(() => ({
@@ -10,7 +12,9 @@ const mocks = vi.hoisted(() => ({
1012
approval: vi.fn(),
1113
projection: vi.fn(),
1214
audit: vi.fn(),
15+
liveAccounts: vi.fn(),
1316
}))
17+
vi.mock('@/lib/sim-search/live/accounts', () => ({ listLiveAccounts: mocks.liveAccounts }))
1418
vi.mock('@/lib/core/application', () => ({ recordProjectedUseCaseAuditEntries: mocks.audit }))
1519
vi.mock('@/lib/core/application/organization-authorization', () => ({
1620
authorizeOrganizationOperation: mocks.authorize,
@@ -81,6 +85,8 @@ const liveBinding = {
8185
describe('organization personal token authorization', () => {
8286
beforeEach(() => {
8387
vi.resetAllMocks()
88+
resetEnvFlagsMock()
89+
mocks.liveAccounts.mockResolvedValue([])
8490
mocks.authorize.mockResolvedValue({ organizationId: 'org', userId: 'person', role: 'member' })
8591
mocks.binding.mockResolvedValue(liveBinding)
8692
mocks.projection.mockReturnValue({ tools: [{ toolId: 'drive_list' }] })
@@ -93,6 +99,36 @@ describe('organization personal token authorization', () => {
9399
mocks.token.mockResolvedValue({ accessToken: 'secret', refreshed: false })
94100
})
95101

102+
it('uses current personal OAuth inventory in live mode without consulting indexed sources', async () => {
103+
setEnvFlags({ isLiveEnterpriseSearchEnabled: true })
104+
mocks.liveAccounts.mockResolvedValue([
105+
{ id: 'own', providerId: 'google-drive', type: 'managed_oauth' },
106+
])
107+
await expect(
108+
resolveOrganizationPersonalToken.execute({ principal, input })
109+
).resolves.toHaveProperty('accessToken', 'secret')
110+
expect(mocks.inventory).not.toHaveBeenCalled()
111+
expect(mocks.liveAccounts).toHaveBeenCalledWith({ organizationId: 'org' }, 'person')
112+
})
113+
it.each(['admin_source', 'service_account', 'personal_token', 'managed_mcp'])(
114+
'never substitutes a %s for a personal OAuth account',
115+
async (type) => {
116+
setEnvFlags({ isLiveEnterpriseSearchEnabled: true })
117+
mocks.liveAccounts.mockResolvedValue([{ id: 'own', providerId: 'google-drive', type }])
118+
await expect(resolveOrganizationPersonalToken.execute({ principal, input })).rejects.toThrow(
119+
'own connected account'
120+
)
121+
expect(mocks.token).not.toHaveBeenCalled()
122+
expect(mocks.inventory).not.toHaveBeenCalled()
123+
}
124+
)
125+
it('observes a revoked live account without falling back to old indexing membership', async () => {
126+
setEnvFlags({ isLiveEnterpriseSearchEnabled: true })
127+
await expect(resolveOrganizationPersonalToken.execute({ principal, input })).rejects.toThrow(
128+
'own connected account'
129+
)
130+
expect(mocks.token).not.toHaveBeenCalled()
131+
})
96132
it('uses the authenticated person inventory and organization token scope without a workspace', async () => {
97133
await expect(resolveOrganizationPersonalToken.execute({ principal, input })).resolves.toEqual({
98134
accessToken: 'secret',

‎apps/sim/lib/credentials/application/resolve-organization-personal-token.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,16 @@ export const resolveOrganizationPersonalToken = {
8383
throw new OrchestrationError('forbidden', 'This integration operation is unavailable.')
8484
}
8585
let cursor: string | undefined
86-
let owned = false
86+
let owned = isLiveEnterpriseSearchEnabled
87+
? (await listLiveAccounts({ organizationId: context.organizationId }, context.userId)).some(
88+
(account) =>
89+
account.id === input.credentialId &&
90+
account.type === 'managed_oauth' &&
91+
account.providerId === binding.providerId
92+
)
93+
: false
8794
const seen = new Set<string>()
88-
for (let page = 0; page < 100; page++) {
95+
for (let page = 0; !isLiveEnterpriseSearchEnabled && page < 100; page++) {
8996
const inventory = await listPersonalSearchIntegrations.execute({
9097
principal,
9198
input: {

‎apps/sim/lib/mothership/assistant/connected-account-tool.test.ts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/** @vitest-environment node */
2+
import { describe, expect, it } from 'vitest'
3+
import {
4+
assistantConnectedAccountTokenParam,
5+
projectAssistantConnectedAccountTool,
6+
} from '@/lib/mothership/assistant/connected-account-tool'
7+
import { getIssueV2Tool } from '@/tools/github/get_issue'
8+
import { searchIssuesV2Tool } from '@/tools/github/search_issues'
9+
10+
describe('GitHub Assistant connected-account adapter', () => {
11+
it('preserves Build and flag-off tool configuration without mutating the registry', () => {
12+
expect(projectAssistantConnectedAccountTool(getIssueV2Tool, false)).toBe(getIssueV2Tool)
13+
const adapted = projectAssistantConnectedAccountTool(getIssueV2Tool, true)
14+
expect(adapted).not.toBe(getIssueV2Tool)
15+
expect(getIssueV2Tool.params.apiKey.required).toBe(true)
16+
expect(getIssueV2Tool.oauth).toBeUndefined()
17+
expect(adapted.oauth).toMatchObject({
18+
required: true,
19+
provider: 'github-repositories',
20+
credentialKind: 'oauth',
21+
})
22+
expect(adapted.params.apiKey).toMatchObject({ visibility: 'hidden', required: false })
23+
expect(adapted.params.credentialId.required).toBe(true)
24+
expect(assistantConnectedAccountTokenParam(adapted)).toBe('apiKey')
25+
})
26+
it('uses the same adapter for the existing issue/PR search tool with total_count', () => {
27+
const adapted = projectAssistantConnectedAccountTool(searchIssuesV2Tool, true)
28+
expect(adapted.request).toBe(searchIssuesV2Tool.request)
29+
expect(adapted.transformResponse).toBe(searchIssuesV2Tool.transformResponse)
30+
expect(adapted.outputs).toBe(searchIssuesV2Tool.outputs)
31+
expect(assistantConnectedAccountTokenParam(adapted)).toBe('apiKey')
32+
})
33+
it('does not turn unrelated API-key tools or GitLab admin sources into personal credentials', () => {
34+
const tool = { ...getIssueV2Tool, id: 'gitlab_get_project' }
35+
expect(projectAssistantConnectedAccountTool(tool, true)).toBe(tool)
36+
expect(assistantConnectedAccountTokenParam(tool)).toBeUndefined()
37+
})
38+
})

‎apps/sim/lib/mothership/assistant/connected-account-tool.ts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import type { ToolMetadata } from '@/tools/metadata'
2+
3+
/** Adapt legacy GitHub API-token operations only at the Assistant boundary; Build schemas stay unchanged. */
4+
export function projectAssistantConnectedAccountTool<T extends ToolMetadata>(
5+
tool: T,
6+
liveSearch: boolean
7+
): T {
8+
if (!liveSearch || !/^github_[a-z0-9_]+$/.test(tool.id) || !tool.params.apiKey || tool.oauth)
9+
return tool
10+
return {
11+
...tool,
12+
oauth: {
13+
required: true,
14+
provider: 'github-repositories',
15+
credentialKind: 'oauth',
16+
requiredScopes: ['repo'],
17+
},
18+
params: {
19+
...tool.params,
20+
apiKey: { ...tool.params.apiKey, required: false, visibility: 'hidden' },
21+
credentialId: {
22+
type: 'string',
23+
required: true,
24+
visibility: 'user-or-llm',
25+
description: 'ID of your connected GitHub account. Authentication is supplied securely.',
26+
},
27+
},
28+
}
29+
}
30+
31+
/** The destination is fixed by trusted registry metadata, never supplied by the model. */
32+
export function assistantConnectedAccountTokenParam(tool: ToolMetadata): 'apiKey' | undefined {
33+
return /^github_[a-z0-9_]+$/.test(tool.id) &&
34+
tool.oauth?.provider === 'github-repositories' &&
35+
tool.params.apiKey?.visibility === 'hidden'
36+
? 'apiKey'
37+
: undefined
38+
}

‎apps/sim/lib/mothership/assistant/tool-policy.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags'
2+
import { projectAssistantConnectedAccountTool } from '@/lib/mothership/assistant/connected-account-tool'
13
import type { ToolMetadata } from '@/tools/metadata'
24

35
export const ASSISTANT_TOOLS = new Set([
@@ -14,6 +16,7 @@ const CREDENTIAL_PARAMS = new Set(['credential', 'credentialId', 'oauthCredentia
1416
/** Assistant uses the regular integration registry, with authentication supplied by the caller's account. */
1517
export function isAssistantIntegrationTool(tool: ToolMetadata | undefined): boolean {
1618
if (!tool) return false
19+
tool = projectAssistantConnectedAccountTool(tool, isLiveEnterpriseSearchEnabled)
1720
const tokenBinding = tool.personalToken
1821
const supportsToken =
1922
tokenBinding && tool.params[tokenBinding.tokenParam] && tool.params[tokenBinding.hostParam]
@@ -34,6 +37,7 @@ export function isAssistantIntegrationTool(tool: ToolMetadata | undefined): bool
3437
}
3538

3639
export function isAssistantIntegrationParameter(tool: ToolMetadata, name: string): boolean {
40+
tool = projectAssistantConnectedAccountTool(tool, isLiveEnterpriseSearchEnabled)
3741
if (CREDENTIAL_PARAMS.has(name)) return true
3842
if (name === tool.personalToken?.tokenParam || name === tool.personalToken?.hostParam)
3943
return false

‎apps/sim/lib/mothership/chat/payload.test.ts‎

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
/**
22
* @vitest-environment node
33
*/
4-
import { envFlagsMockFns, resetEnvFlagsMock, workflowsUtilsMock } from '@sim/testing'
4+
import { envFlagsMockFns, resetEnvFlagsMock, setEnvFlags, workflowsUtilsMock } from '@sim/testing'
55
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
import { getExposedIntegrationTools } from '@/lib/integrations/tool-catalog'
67
import { ChatPayloadSchema } from '@/lib/mothership/generated/protocol'
8+
import { searchIssuesV2Tool } from '@/tools/github/search_issues'
79

810
const {
911
mockCreateUserToolSchema,
@@ -149,13 +151,15 @@ vi.mock('@/tools/params', () => ({
149151

150152
vi.mock('@/tools/metadata', () => ({
151153
getToolMetadata: (id: string) =>
152-
id === 'gmail_send'
153-
? {
154-
id,
155-
params: { accessToken: { type: 'string', visibility: 'hidden', required: true } },
156-
oauth: { required: true, provider: 'google-email' },
157-
}
158-
: undefined,
154+
id === 'github_search_issues_v2'
155+
? searchIssuesV2Tool
156+
: id === 'gmail_send'
157+
? {
158+
id,
159+
params: { accessToken: { type: 'string', visibility: 'hidden', required: true } },
160+
oauth: { required: true, provider: 'google-email' },
161+
}
162+
: undefined,
159163
}))
160164

161165
vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({
@@ -635,6 +639,43 @@ describe('Assistant payload', () => {
635639
mockCreateUserToolSchema.mockReturnValue({ type: 'object', properties: {} })
636640
mockSearchApprovals.mockResolvedValue(new Map())
637641
})
642+
it('discovers the existing GitHub PR-count tool with a personal credential in live Search', async () => {
643+
clearIntegrationToolSchemaCacheForTests()
644+
setEnvFlags({ isLiveEnterpriseSearchEnabled: true })
645+
mockSearchApprovals.mockResolvedValue(new Map([['github', true]]))
646+
vi.mocked(getExposedIntegrationTools).mockReturnValueOnce([
647+
{
648+
toolId: searchIssuesV2Tool.id,
649+
config: searchIssuesV2Tool,
650+
service: 'github',
651+
operation: 'search_issues',
652+
blockType: 'github_v2',
653+
owners: [{ service: 'github', blockType: 'github_v2' }],
654+
},
655+
])
656+
const actualParams = await vi.importActual<typeof import('@/tools/params')>('@/tools/params')
657+
mockCreateUserToolSchema.mockImplementationOnce(actualParams.createUserToolSchema)
658+
const tools = await buildIntegrationToolSchemas('github-live-person', {
659+
schemaSurface: 'copilot',
660+
personalAccountsOnly: true,
661+
organizationId: 'org',
662+
})
663+
expect(tools).toHaveLength(1)
664+
expect(tools[0]).toMatchObject({
665+
name: 'github_search_issues_v2',
666+
oauth: { provider: 'github-repositories' },
667+
input_schema: { required: expect.arrayContaining(['q', 'credentialId']) },
668+
})
669+
expect(tools[0].input_schema.properties).not.toHaveProperty('apiKey')
670+
mockSearchApprovals.mockResolvedValue(new Map([['github', false]]))
671+
expect(
672+
await buildIntegrationToolSchemas('github-live-person', {
673+
schemaSurface: 'copilot',
674+
personalAccountsOnly: true,
675+
organizationId: 'org',
676+
})
677+
).toEqual([])
678+
})
638679
it('advertises approved personal organization integrations and rechecks revocation', async () => {
639680
mockSearchApprovals.mockResolvedValue(new Map([['gmail', true]]))
640681
const options = {

‎apps/sim/lib/mothership/chat/payload.ts‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { LRUCache } from 'lru-cache'
77
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
88
import { isPaid } from '@/lib/billing/plan-helpers'
99
import type { BlockVisibilityState } from '@/lib/core/config/block-visibility'
10-
import { isHosted } from '@/lib/core/config/env-flags'
10+
import { isHosted, isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags'
1111
import { isOAuthServiceDeploymentAvailable } from '@/lib/integrations/availability.server'
1212
import {
1313
type IntegrationGateConfig,
@@ -16,6 +16,7 @@ import {
1616
} from '@/lib/integrations/tool-projection'
1717
import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters'
1818
import { listOrganizationSearchApprovals } from '@/lib/knowledge/search/integration-policy'
19+
import { projectAssistantConnectedAccountTool } from '@/lib/mothership/assistant/connected-account-tool'
1920
import {
2021
isAssistantIntegrationParameter,
2122
isAssistantIntegrationTool,
@@ -212,7 +213,10 @@ export async function buildIntegrationToolSchemas(
212213
)
213214
return structuredClone(
214215
schemas.filter((schema) => {
215-
const metadata = getToolMetadata(schema.name)
216+
const original = getToolMetadata(schema.name)
217+
const metadata = original
218+
? projectAssistantConnectedAccountTool(original, isLiveEnterpriseSearchEnabled)
219+
: undefined
216220
return (
217221
!metadata?.personalToken &&
218222
metadata?.oauth?.required &&
@@ -241,7 +245,10 @@ async function buildIntegrationToolSchemasUncached({
241245
for (const { toolId, config: toolConfig, service, operation } of exposedTools) {
242246
const metadata = getToolMetadata(toolId)
243247
if (options.personalAccountsOnly && !isAssistantIntegrationTool(metadata)) continue
244-
const userSchema = createUserToolSchema(toolConfig, {
248+
const projectedTool = options.personalAccountsOnly
249+
? projectAssistantConnectedAccountTool(toolConfig, isLiveEnterpriseSearchEnabled)
250+
: toolConfig
251+
const userSchema = createUserToolSchema(projectedTool, {
245252
surface: options.schemaSurface,
246253
// On hosted deployments the executor injects hosted keys server-side,
247254
// so the gateway schema must not force the model to supply one (the
@@ -286,11 +293,11 @@ async function buildIntegrationToolSchemasUncached({
286293
}),
287294
defer_loading: true,
288295
executeLocally: catalogEntry?.clientExecutable === true || catalogEntry?.route === 'client',
289-
...(toolConfig.oauth?.required &&
290-
isOAuthServiceDeploymentAvailable(toolConfig.oauth.provider) && {
296+
...(projectedTool.oauth?.required &&
297+
isOAuthServiceDeploymentAvailable(projectedTool.oauth.provider) && {
291298
oauth: {
292299
required: true,
293-
provider: toolConfig.oauth.provider,
300+
provider: projectedTool.oauth.provider,
294301
},
295302
}),
296303
})

0 commit comments

Comments
 (0)