diff --git a/apps/sim/app/api/knowledge/sim-search/personal-integrations/route.ts b/apps/sim/app/api/knowledge/sim-search/personal-integrations/route.ts new file mode 100644 index 00000000000..bdc9d9b6a25 --- /dev/null +++ b/apps/sim/app/api/knowledge/sim-search/personal-integrations/route.ts @@ -0,0 +1,37 @@ +import { + connectPersonalSearchIntegrationContract, + listPersonalSearchIntegrationsContract, +} from '@/lib/api/contracts/knowledge/personal-integrations' +import { + defineInternalJsonRoute, + internalRateLimits, + internalSessionAuth, +} from '@/lib/api/server/routes' +import { internalKnowledgeErrorPolicies } from '@/lib/knowledge/api/route-policies' +import { connectPersonalSearchIntegration } from '@/lib/knowledge/application/connect-personal-search-integration' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { listPersonalSearchIntegrations } from '@/lib/knowledge/application/personal-search-integrations' + +export const GET = defineInternalJsonRoute({ + contract: listPersonalSearchIntegrationsContract, + auth: internalSessionAuth, + operation: knowledgeOperations.listPersonalSearchIntegrations, + rateLimit: internalRateLimits.user({ bucketName: 'knowledge.search.personal-integrations.list' }), + errorPolicy: internalKnowledgeErrorPolicies.connectors, + mapInput: ({ query }) => query, + useCase: listPersonalSearchIntegrations, + present: (data) => ({ success: true as const, data }), +}) + +export const POST = defineInternalJsonRoute({ + contract: connectPersonalSearchIntegrationContract, + auth: internalSessionAuth, + operation: knowledgeOperations.connectPersonalSearchIntegration, + rateLimit: internalRateLimits.user({ + bucketName: 'knowledge.search.personal-integrations.connect', + }), + errorPolicy: internalKnowledgeErrorPolicies.connectAccount, + mapInput: ({ body }) => body, + useCase: connectPersonalSearchIntegration, + present: (data) => ({ success: true as const, data }), +}) diff --git a/apps/sim/app/o/[organizationId]/home/organization-home.tsx b/apps/sim/app/o/[organizationId]/home/organization-home.tsx index 8d4ae29eb82..59d754b9117 100644 --- a/apps/sim/app/o/[organizationId]/home/organization-home.tsx +++ b/apps/sim/app/o/[organizationId]/home/organization-home.tsx @@ -6,6 +6,7 @@ import { MothershipHandoffStorage } from '@/lib/core/utils/browser-storage' import { Composer } from '@/app/o/[organizationId]/home/components/composer' import { GetStarted } from '@/app/o/[organizationId]/home/components/get-started' import { useOrganizationContext } from '@/app/o/[organizationId]/providers/organization-provider' +import { SearchIntegrationConnection } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection' import { MothershipChat } from '@/app/workspace/[workspaceId]/home/components/mothership-chat' import { useChat } from '@/app/workspace/[workspaceId]/home/hooks/use-chat' import { useMarkMothershipChatRead } from '@/hooks/queries/mothership-chats' @@ -79,6 +80,7 @@ function OrganizationHomeContent({ userName, chatId }: OrganizationHomeProps) {
{hasChat ? ( + {connectionRequest && ( + + )}
{sources.isError && !sources.isFetchNextPageError ? ( { }) describe('integrations page Slack context', () => { + it('preserves a requested connection across login and validates it in the existing organization page', async () => { + const selected = { + ...props, + searchParams: Promise.resolve({ + connectorType: 'gmail', + connectorId: 'source', + credentialId: 'account', + }), + } + const page = await OrganizationIntegrationsPage(selected) + expect(page.props.connectionRequest).toMatchObject({ + userId: 'viewer', + target: { + type: 'link', + connectorType: 'gmail', + connectorId: 'source', + credentialId: 'account', + }, + }) + authMockFns.mockGetSession.mockResolvedValue(null) + await expect(OrganizationIntegrationsPage(selected)).rejects.toThrow('Redirect') + expect(mocks.redirect).toHaveBeenCalledWith( + `/login?callbackUrl=${encodeURIComponent('/o/organization-a/integrations?connectorType=gmail&connectorId=source&credentialId=account')}` + ) + }) + it('rejects unknown providers and reconnects without a source', async () => { + for (const query of [ + { connectorType: 'invented' }, + { connectorType: 'gmail', credentialId: 'account' }, + ]) { + await expect( + OrganizationIntegrationsPage({ ...props, searchParams: Promise.resolve(query) }) + ).rejects.toThrow('Not found') + } + }) it('preserves the source page and Slack question context through login', async () => { authMockFns.mockGetSession.mockResolvedValue(null) await expect(OrganizationIntegrationsPage(props)).rejects.toThrow('Redirect') diff --git a/apps/sim/app/o/[organizationId]/integrations/page.tsx b/apps/sim/app/o/[organizationId]/integrations/page.tsx index 23d81a29386..e4bc2ddcf87 100644 --- a/apps/sim/app/o/[organizationId]/integrations/page.tsx +++ b/apps/sim/app/o/[organizationId]/integrations/page.tsx @@ -2,11 +2,17 @@ import type { Metadata } from 'next' import { notFound, redirect } from 'next/navigation' import { slackSearchOnboardingInputSchema } from '@/lib/api/contracts/knowledge/slack' import { getSession } from '@/lib/auth' +import { + searchConnectionPath, + searchConnectionTargetSchema, +} from '@/lib/knowledge/search/connection-target' import { organizationRoutes } from '@/lib/navigation/paths' import { getOrganizationSurfaceContext } from '@/lib/organizations/surface' +import { SEARCH_CONNECTORS } from '@/lib/sim-search/connectors' import { slackSearchIntegrationsPath } from '@/lib/slack-search/onboarding' import { buildAuthCrossLink } from '@/app/(auth)/auth-redirect' import { OrganizationIntegrations } from '@/app/o/[organizationId]/integrations/integrations' +import { loadIntegrationConnectionParams } from '@/app/o/[organizationId]/integrations/search-params' export const metadata: Metadata = { title: 'Integrations', @@ -15,7 +21,12 @@ export const metadata: Metadata = { interface OrganizationIntegrationsPageProps { params: Promise<{ organizationId: string }> - searchParams: Promise<{ slack?: string | string[] }> + searchParams: Promise<{ + slack?: string | string[] + connectorType?: string | string[] + connectorId?: string | string[] + credentialId?: string | string[] + }> } export default async function OrganizationIntegrationsPage({ @@ -23,7 +34,22 @@ export default async function OrganizationIntegrationsPage({ searchParams, }: OrganizationIntegrationsPageProps) { const { organizationId } = await params - const { slack } = await searchParams + const query = await searchParams + const { slack } = query + const selection = loadIntegrationConnectionParams(query) + const connector = SEARCH_CONNECTORS.find((entry) => entry.type === selection.connectorType) + const requested = + selection.connectorType || selection.connectorId || selection.credentialId + ? searchConnectionTargetSchema.safeParse({ + type: 'link', + provider: connector?.providerId, + connectorType: selection.connectorType, + ...(selection.connectorId ? { connectorId: selection.connectorId } : {}), + ...(selection.credentialId ? { credentialId: selection.credentialId } : {}), + }) + : undefined + if (requested && !requested.success) notFound() + const connectionTarget = requested?.data const context = slack === undefined ? undefined : slackSearchOnboardingInputSchema.safeParse({ token: slack }) if (context && !context.success) notFound() @@ -32,9 +58,11 @@ export default async function OrganizationIntegrationsPage({ if (!session?.user) redirect( buildAuthCrossLink('/login', { - callbackUrl: slackToken - ? slackSearchIntegrationsPath(organizationId, slackToken) - : organizationRoutes(organizationId).integrations, + callbackUrl: connectionTarget + ? searchConnectionPath(organizationId, connectionTarget) + : slackToken + ? slackSearchIntegrationsPath(organizationId, slackToken) + : organizationRoutes(organizationId).integrations, isInviteFlow: false, }) ) @@ -42,6 +70,9 @@ export default async function OrganizationIntegrationsPage({ if (!organizationContext?.searchAccess.memberScoped) notFound() return ( ) diff --git a/apps/sim/app/o/[organizationId]/integrations/search-params.ts b/apps/sim/app/o/[organizationId]/integrations/search-params.ts new file mode 100644 index 00000000000..55ccd0e2a41 --- /dev/null +++ b/apps/sim/app/o/[organizationId]/integrations/search-params.ts @@ -0,0 +1,9 @@ +import { createLoader, parseAsString } from 'nuqs/server' + +export const integrationConnectionParams = { + connectorType: parseAsString.withDefault(''), + connectorId: parseAsString.withDefault(''), + credentialId: parseAsString.withDefault(''), +} + +export const loadIntegrationConnectionParams = createLoader(integrationConnectionParams) diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/chat-surface-context/chat-surface-context.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/chat-surface-context/chat-surface-context.tsx index 6864df15943..f4d8cd5954e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/chat-surface-context/chat-surface-context.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/chat-surface-context/chat-surface-context.tsx @@ -1,6 +1,7 @@ 'use client' import { + type ComponentType, createContext, type ReactNode, useCallback, @@ -10,6 +11,7 @@ import { useRef, } from 'react' import { noop } from '@sim/utils/helpers' +import type { SearchIntegrationConnectionProps } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection' import type { WorkspaceResourceRef } from '@/app/workspace/[workspaceId]/home/types' import type { ChatContext } from '@/stores/panel' @@ -20,6 +22,7 @@ import type { ChatContext } from '@/stores/panel' * consume them without relaying through every intermediate component. */ interface ChatSurfaceContextValue { + SearchConnectionComponent?: ComponentType /** Resolved id of the chat backing this surface, if one exists yet. */ chatId?: string /** Id of the user interacting with this surface. */ @@ -44,6 +47,7 @@ const ChatSurfaceContext = createContext({ }) interface ChatSurfaceProviderProps { + SearchConnectionComponent?: ComponentType chatId?: string userId?: string onContextAdd?: (context: ChatContext) => void @@ -59,6 +63,7 @@ interface ChatSurfaceProviderProps { * not re-render when a parent re-creates a handler. */ export function ChatSurfaceProvider({ + SearchConnectionComponent, chatId, userId, onContextAdd, @@ -88,13 +93,21 @@ export function ChatSurfaceProvider({ const value = useMemo( () => ({ + SearchConnectionComponent, chatId, userId, onContextAdd: stableOnContextAdd, onContextRemove: stableOnContextRemove, onWorkspaceResourceSelect: stableOnWorkspaceResourceSelect, }), - [chatId, userId, stableOnContextAdd, stableOnContextRemove, stableOnWorkspaceResourceSelect] + [ + SearchConnectionComponent, + chatId, + userId, + stableOnContextAdd, + stableOnContextRemove, + stableOnWorkspaceResourceSelect, + ] ) return {children} diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/interaction-card.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/interaction-card.tsx index 31ac390e5bf..9fe74690780 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/interaction-card.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/interaction-card.tsx @@ -90,6 +90,8 @@ export const InteractionCardInputRow = forwardRef @@ -123,7 +128,7 @@ export function InteractionCardActionRow({ > {label} - + {trailing ?? } ) } diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection.tsx new file mode 100644 index 00000000000..e1df655f9ce --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection.tsx @@ -0,0 +1,117 @@ +'use client' + +import { useState } from 'react' +import { Chip } from '@sim/emcn' +import { Check } from '@sim/emcn/icons' +import type { SearchConnectionTarget } from '@/lib/knowledge/search/connection-target' +import { SEARCH_CONNECTORS } from '@/lib/sim-search/connectors' +import { + InteractionCard, + InteractionCardActionRow, +} from '@/app/workspace/[workspaceId]/home/components/message-content/components/interaction-card' +import { SourceSetupModal } from '@/app/workspace/[workspaceId]/home/components/search-sources/source-setup-modal' +import { BrandIcon } from '@/blocks/brand-icon' +import { useSearchIntegrationConnection } from '@/hooks/use-search-integration-connection' + +export interface SearchIntegrationConnectionProps { + organizationId: string + userId: string + target: SearchConnectionTarget + controlId: string + embedded?: boolean + divided?: boolean + onConnected?: () => void +} + +/** The ordinary credential card, with personal Search enrollment instead of workspace credentials. */ +export function SearchIntegrationConnection(props: SearchIntegrationConnectionProps) { + return ( + + ) +} + +function SearchIntegrationConnectionControl({ + organizationId, + userId, + target, + controlId, + embedded, + divided, + onConnected, +}: SearchIntegrationConnectionProps) { + const [setupOpen, setSetupOpen] = useState(false) + const connector = SEARCH_CONNECTORS.find((entry) => entry.type === target.connectorType) + const connection = useSearchIntegrationConnection({ + organizationId, + userId, + target, + controlId, + onConnected, + }) + const name = connector?.meta.name ?? target.provider + const action = target.credentialId ? 'Reconnect' : 'Connect' + const label = connection.connected + ? `Connected ${name}` + : connection.isLoading + ? `Checking ${name} connections…` + : connection.pending + ? `Waiting for ${name} connection…` + : !connection.available + ? `${name} connection is no longer available` + : `${action} ${name}` + const handleConnect = () => { + if (connector && !connection.connectorId && connector.setupFields.length && !connection.pending) + setSetupOpen(true) + else void connection.connect() + } + const content = ( + <> + + } + trailing={ + connection.connected ? ( + + ) : undefined + } + /> + {connection.pending && Cancel} + {connection.error && ( +
+ {connection.error}{' '} + void connection.retry() : handleConnect}> + Retry + +
+ )} + {setupOpen && connector && ( + setSetupOpen(false)} + isPending={connection.isStarting} + error={connection.error} + onConnect={(config) => { + void connection.connect(config).then((started) => { + if (started) setSetupOpen(false) + }) + }} + /> + )} + + ) + return embedded ? content : {content} +} diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.tsx index daa8df37bf0..797686b7699 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.tsx @@ -7,6 +7,7 @@ import { createRoot, type Root } from 'react-dom/client' import { beforeEach, describe, expect, it, vi } from 'vitest' const { + mockParams, mockRefetchPersonalEnvironment, mockRefetchWorkspaceCredentials, mockIsBrowserAgentAvailable, @@ -18,6 +19,7 @@ const { mockUseWorkspaceCredential, mockUseWorkspaceCredentials, } = vi.hoisted(() => ({ + mockParams: vi.fn(() => ({ workspaceId: 'workspace-1' })), mockUpdateWorkspaceCredential: vi.fn(async () => undefined), mockRefetchPersonalEnvironment: vi.fn(async () => ({ data: {} })), mockRefetchWorkspaceCredentials: vi.fn(async () => ({ data: [] })), @@ -35,7 +37,20 @@ vi.mock('@/app/workspace/[workspaceId]/providers/workspace-permissions-provider' })) vi.mock('next/navigation', () => ({ - useParams: () => ({ workspaceId: 'workspace-1' }), + useParams: mockParams, +})) + +vi.mock('@/lib/auth/auth-client', () => ({ + useSession: () => ({ data: { user: { id: 'person' } } }), +})) +vi.mock('@/app/workspace/[workspaceId]/home/components/chat-surface-context', () => ({ + useChatSurface: () => ({ + SearchConnectionComponent: ({ onConnected }: { onConnected?: () => void }) => ( + + ), + }), })) vi.mock('@/hooks/queries/credentials', () => ({ @@ -98,6 +113,7 @@ describe('CredentialDisplay link tag', () => { beforeEach(() => { ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true vi.clearAllMocks() + mockParams.mockReturnValue({ workspaceId: 'workspace-1' }) window.localStorage.clear() window.history.replaceState({}, '', '/workspace/workspace-1/chat/chat-1') mockUseUserPermissionsContext.mockReturnValue({ canEdit: true }) @@ -110,6 +126,38 @@ describe('CredentialDisplay link tag', () => { mockIsBrowserAgentAvailable.mockReturnValue(false) }) + it('keeps organization Search connection completion behind Submit', async () => { + mockParams.mockReturnValue({ organizationId: 'org' } as never) + const container = document.createElement('div') + const root = createRoot(container) + const onContinue = vi.fn() + act(() => + root.render( + + ) + ) + const connect = Array.from(container.querySelectorAll('button')).find( + (button) => button.textContent === 'Test Search connection' + ) + act(() => connect?.click()) + expect(onContinue).not.toHaveBeenCalled() + const submit = Array.from(container.querySelectorAll('button')).find( + (button) => button.textContent === 'Submit' + ) + expect(submit).toBeDefined() + await act(async () => submit?.click()) + expect(onContinue).toHaveBeenCalledOnce() + expect(onContinue.mock.calls[0][0]).toContain('connected') + act(() => root.unmount()) + }) it('renders browser takeover through the shared question UI', () => { mockIsBrowserAgentAvailable.mockReturnValue(true) const { container, root } = renderCredentialLink({ diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx index 0330603d98c..889b2a15353 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx @@ -19,6 +19,14 @@ import { resolveOAuthServiceForSlug, resolveServiceAccountIntegration, } from '@/lib/integrations/oauth-service' +import { + readSearchConnectionAttempt, + searchConnectionAttemptKey, +} from '@/lib/knowledge/search/connection-attempt' +import { + parseSearchConnectionBody, + searchConnectionTargetSchema, +} from '@/lib/knowledge/search/connection-target' import { OAUTH_PROVIDERS } from '@/lib/oauth/oauth' import { getServiceConfigByProviderId } from '@/lib/oauth/utils' import { finishTerminalHandoff, isTerminalAvailable } from '@/lib/terminal/transport' @@ -150,6 +158,9 @@ export interface CredentialItemData { * rotate the secret on this credential; absent = create a new one. */ credentialId?: string + /** Canonical Search source requested by an organization connection control. */ + connectorType?: string + connectorId?: string } /** @@ -520,6 +531,8 @@ function isCredentialItemData(value: unknown): value is CredentialItemData { } return typeof value.provider === 'string' && value.provider.trim().length > 0 } + if (value.type === 'link' && value.connectorType !== undefined) + return searchConnectionTargetSchema.safeParse(value).success if (value.type === 'link' && value.value === undefined) { return typeof value.provider === 'string' && value.provider.trim().length > 0 } @@ -538,6 +551,8 @@ export function parseCredentialTagBody(body: string): CredentialTagData | null { try { const parsed = JSON.parse(body) as unknown const items = Array.isArray(parsed) ? parsed : [parsed] + if (items.some((item) => isRecordLike(item) && item.connectorType !== undefined)) + return parseSearchConnectionBody(body) return items.length > 0 && items.every(isCredentialItemData) ? items : null } catch { return null @@ -2691,7 +2706,7 @@ export function credentialTagHasVisibleCard( function CredentialItemDisplay({ data, requestMode, - controlId, + controlId = 'credential-link', embedded = false, divided = false, secretValue, @@ -2699,6 +2714,9 @@ function CredentialItemDisplay({ onSaved, onConnected, }: CredentialControlProps) { + const { organizationId } = useParams<{ organizationId?: string }>() + const { SearchConnectionComponent } = useChatSurface() + const { data: session } = useSession() if ( requestMode === 'assistant' && data.type !== 'link' && @@ -2738,6 +2756,23 @@ function CredentialItemDisplay({ if (data.type === 'link') { if (requestMode === 'assistant') { + if (organizationId) { + const target = searchConnectionTargetSchema.safeParse(data) + if (!target.success || !session?.user?.id) return null + if (!SearchConnectionComponent) + throw new Error('Search connection controls require an organization chat surface') + return ( + + ) + } return ( void }) { - const { workspaceId } = useParams<{ workspaceId: string }>() + const { workspaceId, organizationId } = useParams<{ + workspaceId: string + organizationId?: string + }>() + const { data: session } = useSession() const { canEdit } = useUserPermissionsContext() const upsertWorkspace = useUpsertWorkspaceEnvironment() const savePersonal = useSavePersonalEnvironment() @@ -2829,6 +2868,19 @@ function CredentialInputCard({ if (item.type !== 'link' && item.type !== 'service_account') continue const index = restoreIndex++ if (item.type !== 'link') continue + if (requestMode === 'assistant' && organizationId && session?.user?.id) { + const target = searchConnectionTargetSchema.safeParse(item) + if (!target.success) continue + const attempt = readSearchConnectionAttempt( + searchConnectionAttemptKey( + organizationId, + session.user.id, + `${controlIdPrefix}:${dataIndex}:${JSON.stringify(target.data)}` + ) + ) + if (attempt?.status === 'connected') restored.add(index) + continue + } const { providerId, reconnectCredentialId } = requestMode === 'assistant' ? { @@ -2851,7 +2903,15 @@ function CredentialInputCard({ if (Array.from(restored).every((index) => current.has(index))) return current return new Set([...current, ...restored]) }) - }, [abandoned, controlIdPrefix, data, workspaceId, requestMode]) + }, [ + abandoned, + controlIdPrefix, + data, + workspaceId, + requestMode, + organizationId, + session?.user?.id, + ]) let integrationIndex = 0 let secretIndex = 0 diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx index d84c76b2685..6fc88b74a04 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx @@ -1,5 +1,6 @@ 'use client' +import type { ComponentType } from 'react' import { memo, type ReactNode, @@ -34,6 +35,7 @@ import { parseLastCredentialTag, parseLastQuestionTag, } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags' +import type { SearchIntegrationConnectionProps } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection' import { prepareCopyableMarkdown, toCopyableMarkdown, @@ -62,6 +64,7 @@ import { MothershipChatSkeleton } from './components/mothership-chat-skeleton' import { shouldShowAssistantMessageActions } from './message-actions-visibility' interface MothershipChatProps { + SearchConnectionComponent?: ComponentType workspaceId?: string composer?: ReactNode messages: ChatMessage[] @@ -333,6 +336,7 @@ const AssistantMessageRow = memo(function AssistantMessageRow({ }) export function MothershipChat({ + SearchConnectionComponent, workspaceId, composer, messages: messagesProp, @@ -781,6 +785,7 @@ export function MothershipChat({ return ( [...personalSearchIntegrationKeys.all, 'list'] as const, + list: (query: PersonalSearchIntegrationsQuery) => + [...personalSearchIntegrationKeys.lists(), query] as const, +} + +export function usePersonalSearchIntegrations( + query: PersonalSearchIntegrationsQuery, + options: { pending?: boolean } = {} +) { + return useQuery({ + queryKey: personalSearchIntegrationKeys.list(query), + queryFn: async ({ signal }) => + (await requestJson(listPersonalSearchIntegrationsContract, { query, signal })).data, + staleTime: PERSONAL_SEARCH_INTEGRATIONS_STALE_TIME, + refetchInterval: options.pending ? 1_500 : false, + }) +} + +export function useConnectPersonalSearchIntegration() { + const client = useQueryClient() + return useMutation({ + mutationFn: async (body: ConnectPersonalSearchIntegrationBody) => + (await requestJson(connectPersonalSearchIntegrationContract, { body })).data, + onSettled: (_data, _error, body) => + Promise.all([ + client.invalidateQueries({ queryKey: personalSearchIntegrationKeys.lists() }), + client.invalidateQueries({ + queryKey: searchSourceKeys.list({ + kind: 'organization', + organizationId: body.organizationId, + }), + }), + client.invalidateQueries({ + queryKey: organizationAccountsKeys.detail(body.organizationId), + }), + ]), + }) +} diff --git a/apps/sim/hooks/use-search-integration-connection.test.tsx b/apps/sim/hooks/use-search-integration-connection.test.tsx new file mode 100644 index 00000000000..421180cab62 --- /dev/null +++ b/apps/sim/hooks/use-search-integration-connection.test.tsx @@ -0,0 +1,239 @@ +/** @vitest-environment jsdom */ +import { act } from 'react' +import { createRoot, type Root } from 'react-dom/client' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const m = vi.hoisted(() => ({ + mutate: vi.fn(), + invalidate: vi.fn(), + refetch: vi.fn(), + connected: vi.fn(), + receipts: new Map(), + accounts: [] as Array<{ credentialId: string; status: string }>, + channels: [] as Array<{ + name: string + onmessage: ((event: MessageEvent) => void) | null + close: () => void + }>, + queryError: null as Error | null, + requestedTarget: undefined as + | { type: 'link'; provider: string; connectorType: string; connectorId?: string } + | undefined, +})) +const target = { + type: 'link', + provider: 'gmail', + connectorType: 'gmail', + connectorId: 'source', +} as const +const client = { invalidateQueries: m.invalidate } +vi.mock('@tanstack/react-query', () => ({ useQueryClient: () => client })) +vi.mock('@/hooks/queries/personal-search-integrations', () => ({ + personalSearchIntegrationKeys: { lists: () => ['personal-integrations', 'list'] }, + useConnectPersonalSearchIntegration: () => ({ mutateAsync: m.mutate, isPending: false }), + usePersonalSearchIntegrations: (query: { completionId?: string }) => ({ + data: { + connections: [{ accounts: m.accounts }], + available: [{ target }], + completedCredentialId: m.receipts.get(query.completionId ?? '') ?? null, + }, + isSuccess: !m.queryError, + isPending: false, + error: m.queryError, + refetch: m.refetch, + }), +})) +vi.mock('@/hooks/queries/organization-accounts', () => ({ + organizationAccountsKeys: { detail: (id: string) => ['accounts', id] }, +})) + +import { useSearchIntegrationConnection } from '@/hooks/use-search-integration-connection' + +type Connection = ReturnType +const latest = new Map() +let root: Root +let container: HTMLDivElement +const windows: Array<{ + location: { href: string } + close: ReturnType + focus: ReturnType + closed: boolean +}> = [] +function Harness({ id }: { id: string }) { + latest.set( + id, + useSearchIntegrationConnection({ + organizationId: 'org', + userId: 'person', + target: m.requestedTarget ?? target, + controlId: id, + onConnected: m.connected, + }) + ) + return null +} +function render(ids = ['one']) { + act(() => + root.render( + <> + {ids.map((id) => ( + + ))} + + ) + ) +} +function connection(id = 'one') { + const value = latest.get(id) + if (!value) throw new Error('Missing connection') + return value +} +beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers() + m.accounts = [] + m.receipts.clear() + m.channels.length = 0 + m.queryError = null + m.requestedTarget = undefined + windows.length = 0 + window.localStorage.clear() + ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true + vi.stubGlobal( + 'BroadcastChannel', + class { + onmessage = null + close = vi.fn() + constructor(public name: string) { + m.channels.push(this) + } + } + ) + vi.spyOn(window, 'open').mockImplementation(() => { + const popup = { location: { href: '' }, close: vi.fn(), focus: vi.fn(), closed: false } + windows.push(popup) + return popup as unknown as Window + }) + m.refetch.mockResolvedValue({ isSuccess: true, data: { connections: [] } }) + m.mutate.mockResolvedValue({ + url: 'https://provider.test/authorize', + connectorId: 'source', + knowledgeBaseId: 'kb', + }) + container = document.createElement('div') + document.body.append(container) + root = createRoot(container) +}) +afterEach(() => { + act(() => root.unmount()) + container.remove() + latest.clear() + vi.useRealTimers() + vi.restoreAllMocks() + vi.unstubAllGlobals() +}) + +describe('Search connection card lifecycle', () => { + it('starts OAuth only on click and completes only when its receipt and current account agree', async () => { + render() + expect(m.mutate).not.toHaveBeenCalled() + await act(async () => { + await connection().connect() + }) + const id = m.mutate.mock.calls[0][0].oauthCompletionId + expect(windows[0].location.href).toBe('https://provider.test/authorize') + expect(connection().pending).toBe(true) + act(() => m.channels[0].onmessage?.(new MessageEvent('message', { data: 'connected' }))) + expect(connection().connected).toBe(false) + m.accounts = [{ credentialId: 'mine', status: 'connected' }] + m.receipts.set(id, 'mine') + render() + expect(connection().connected).toBe(true) + expect(m.connected).toHaveBeenCalled() + }) + it('does not complete another card waiting for the same provider', async () => { + render(['one', 'two']) + await act(async () => { + await connection('one').connect() + await connection('two').connect() + }) + const [first, second] = m.mutate.mock.calls.map(([input]) => input.oauthCompletionId) + expect(first).not.toBe(second) + m.accounts = [{ credentialId: 'mine', status: 'connected' }] + m.receipts.set(first, 'mine') + render(['one', 'two']) + expect(connection('one').connected).toBe(true) + expect(connection('two').connected).toBe(false) + expect(windows[1].close).not.toHaveBeenCalled() + }) + it('restores pending completion on reload and ignores browser clock skew', async () => { + render() + vi.setSystemTime(new Date('2030-01-01')) + await act(async () => { + await connection().connect() + }) + const id = m.mutate.mock.calls[0][0].oauthCompletionId + act(() => root.unmount()) + root = createRoot(container) + m.accounts = [{ credentialId: 'mine', status: 'connected' }] + m.receipts.set(id, 'mine') + render() + expect(connection().connected).toBe(true) + m.accounts = [{ credentialId: 'mine', status: 'reconnect_needed' }] + render() + expect(connection().connected).toBe(false) + }) + it('supports cancel and retry with a fresh attempt', async () => { + render() + await act(async () => { + await connection().connect() + }) + act(() => connection().cancel()) + expect(connection().pending).toBe(false) + expect(connection().error).toContain('canceled') + await act(async () => { + await connection().connect() + }) + expect(m.mutate).toHaveBeenCalledTimes(2) + expect(connection().pending).toBe(true) + }) + it('retries the exact source created by the first attempt after cancellation or reload', async () => { + m.requestedTarget = { type: 'link', provider: 'gmail', connectorType: 'gmail' } + render() + await act(async () => { + await connection().connect() + }) + expect(m.mutate.mock.calls[0][0].target.connectorId).toBeUndefined() + act(() => connection().cancel()) + act(() => root.unmount()) + root = createRoot(container) + render() + expect(connection().available).toBe(true) + await act(async () => { + await connection().connect() + }) + expect(m.mutate.mock.calls[1][0].target.connectorId).toBe('source') + expect(m.mutate.mock.calls[1][0].target.credentialId).toBeUndefined() + }) + it('keeps failed starts actionable and rejects stale data after authorization errors', async () => { + render() + m.mutate.mockRejectedValueOnce(new Error('Source no longer available')) + await act(async () => { + expect(await connection().connect()).toBe(false) + }) + expect(connection().error).toBe('Source no longer available') + expect(windows[0].close).toHaveBeenCalled() + m.queryError = new Error('Membership revoked') + render() + expect(connection().available).toBe(false) + }) + it('times out without closing a separate successful attempt', async () => { + render() + await act(async () => { + await connection().connect() + }) + act(() => vi.advanceTimersByTime(600_001)) + expect(connection().pending).toBe(false) + expect(connection().error).toContain('timed out') + }) +}) diff --git a/apps/sim/hooks/use-search-integration-connection.ts b/apps/sim/hooks/use-search-integration-connection.ts new file mode 100644 index 00000000000..96c455ac16f --- /dev/null +++ b/apps/sim/hooks/use-search-integration-connection.ts @@ -0,0 +1,229 @@ +'use client' + +import { useCallback, useEffect, useRef, useState } from 'react' +import { getErrorMessage } from '@sim/utils/errors' +import { generateId } from '@sim/utils/id' +import { useQueryClient } from '@tanstack/react-query' +import { + CREDENTIAL_GROUP_OAUTH_FAILURE_MESSAGES, + credentialGroupOAuthCompletionChannel, + isCredentialGroupOAuthFailure, +} from '@/lib/credential-groups/oauth-completion' +import { + readSearchConnectionAttempt, + SEARCH_CONNECTION_ATTEMPT_EVENT, + SEARCH_CONNECTION_ATTEMPT_MAX_AGE_MS, + type SearchConnectionAttempt, + searchConnectionAttemptKey, + writeSearchConnectionAttempt, +} from '@/lib/knowledge/search/connection-attempt' +import type { SearchConnectionTarget } from '@/lib/knowledge/search/connection-target' +import { organizationAccountsKeys } from '@/hooks/queries/organization-accounts' +import { + personalSearchIntegrationKeys, + useConnectPersonalSearchIntegration, + usePersonalSearchIntegrations, +} from '@/hooks/queries/personal-search-integrations' +import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys' + +interface SearchIntegrationConnectionProps { + organizationId: string + userId: string + target: SearchConnectionTarget + controlId: string + onConnected?: () => void +} + +/** OAuth stays user initiated; card completion is proved by the current personal inventory. */ +export function useSearchIntegrationConnection({ + organizationId, + userId, + target, + controlId, + onConnected, +}: SearchIntegrationConnectionProps) { + const key = searchConnectionAttemptKey( + organizationId, + userId, + `${controlId}:${JSON.stringify(target)}` + ) + const [attempt, setAttempt] = useState(() => readSearchConnectionAttempt(key)) + const [localError, setLocalError] = useState(null) + const popup = useRef(null) + const starting = useRef(false) + const callback = useRef(onConnected) + useEffect(() => { + callback.current = onConnected + }, [onConnected]) + const client = useQueryClient() + const { mutateAsync, isPending } = useConnectPersonalSearchIntegration() + const pending = attempt?.status === 'pending' + const connectorId = target.connectorId ?? attempt?.connectorId + const effectiveTarget = connectorId ? { ...target, connectorId } : target + const query = usePersonalSearchIntegrations( + { + organizationId, + connectorType: target.connectorType, + connectorId, + completionId: attempt?.completionId, + }, + { pending } + ) + const accounts = query.data?.connections.flatMap((entry) => entry.accounts) ?? [] + const connected = + query.isSuccess && + attempt !== null && + attempt.status !== 'failed' && + accounts.some( + (account) => + account.status === 'connected' && + account.credentialId === + (query.data?.completedCredentialId ?? + (attempt.status === 'connected' ? attempt.credentialId : undefined)) && + (!target.credentialId || account.credentialId === target.credentialId) + ) + const availableTargets = [ + ...(query.data?.available.map((entry) => entry.target) ?? []), + ...(query.data?.connections.flatMap((entry) => + entry.accounts.flatMap((account) => (account.action ? [account.action] : [])) + ) ?? []), + ] + const available = + query.isSuccess && + availableTargets.some( + (candidate) => JSON.stringify(candidate) === JSON.stringify(effectiveTarget) + ) + + useEffect(() => { + const refresh = () => setAttempt(readSearchConnectionAttempt(key)) + window.addEventListener(SEARCH_CONNECTION_ATTEMPT_EVENT, refresh) + window.addEventListener('storage', refresh) + refresh() + return () => { + window.removeEventListener(SEARCH_CONNECTION_ATTEMPT_EVENT, refresh) + window.removeEventListener('storage', refresh) + } + }, [key]) + + useEffect(() => { + if (!attempt || attempt.status !== 'pending') return + const channel = new BroadcastChannel( + credentialGroupOAuthCompletionChannel(attempt.completionId) + ) + const fail = (error: string) => + writeSearchConnectionAttempt(key, { ...attempt, status: 'failed', error }) + channel.onmessage = ({ data }: MessageEvent) => { + if (isCredentialGroupOAuthFailure(data)) fail(CREDENTIAL_GROUP_OAUTH_FAILURE_MESSAGES[data]) + else if (data === 'connected') { + void client.invalidateQueries({ queryKey: personalSearchIntegrationKeys.lists() }) + void client.invalidateQueries({ + queryKey: searchSourceKeys.list({ kind: 'organization', organizationId }), + }) + void client.invalidateQueries({ queryKey: organizationAccountsKeys.detail(organizationId) }) + } + } + const timer = window.setTimeout( + () => { + popup.current?.close() + fail('Connection timed out. Try connecting again.') + }, + Math.max(0, attempt.requestedAt + SEARCH_CONNECTION_ATTEMPT_MAX_AGE_MS - Date.now()) + ) + return () => { + channel.close() + window.clearTimeout(timer) + } + }, [attempt, key, client, organizationId]) + + useEffect(() => { + if (!connected || !attempt) return + if (attempt.status !== 'connected') + writeSearchConnectionAttempt(key, { + ...attempt, + status: 'connected', + error: null, + credentialId: query.data?.completedCredentialId ?? attempt.credentialId, + }) + popup.current?.close() + callback.current?.() + }, [connected, attempt, key, query.data?.completedCredentialId]) + + const { refetch } = query + const connect = useCallback( + async (sourceConfig?: Record) => { + if (starting.current || isPending || connected) return + if (pending && popup.current && !popup.current.closed) { + popup.current.focus() + return + } + const tab = window.open('about:blank', '_blank', 'width=600,height=700') + if (!tab) { + setLocalError('Allow pop-ups for this site to connect your account.') + return + } + tab.opener = null + popup.current = tab + starting.current = true + setLocalError(null) + let next: SearchConnectionAttempt | undefined + try { + const fresh = await refetch() + if (!fresh.isSuccess) throw fresh.error + next = { + completionId: generateId(), + requestedAt: Date.now(), + connectorId, + status: 'pending', + error: null, + } + writeSearchConnectionAttempt(key, next) + const result = await mutateAsync({ + organizationId, + target: connectorId ? { ...target, connectorId } : target, + sourceConfig, + oauthCompletionId: next.completionId, + }) + const url = new URL(result.url) + if ( + url.protocol !== 'https:' && + !(url.protocol === 'http:' && url.origin === window.location.origin) + ) + throw new Error('The provider authorization URL is invalid') + writeSearchConnectionAttempt(key, { ...next, connectorId: result.connectorId }) + tab.location.href = url.href + return true + } catch (error) { + tab.close() + const message = getErrorMessage(error, 'Could not start the connection') + if (next) writeSearchConnectionAttempt(key, { ...next, status: 'failed', error: message }) + setLocalError(message) + return false + } finally { + starting.current = false + } + }, + [isPending, connected, pending, refetch, mutateAsync, organizationId, target, connectorId, key] + ) + const cancel = useCallback(() => { + popup.current?.close() + if (attempt?.status === 'pending') + writeSearchConnectionAttempt(key, { + ...attempt, + status: 'failed', + error: 'Connection canceled. You can try again.', + }) + }, [attempt, key]) + return { + connect, + cancel, + inventoryError: query.error?.message, + connected, + connectorId, + pending, + available, + isStarting: isPending, + isLoading: query.isPending, + error: localError ?? query.error?.message ?? attempt?.error, + retry: query.refetch, + } +} diff --git a/apps/sim/lib/api/contracts/knowledge/connectors.ts b/apps/sim/lib/api/contracts/knowledge/connectors.ts index 438765f3c7b..ce2185c981d 100644 --- a/apps/sim/lib/api/contracts/knowledge/connectors.ts +++ b/apps/sim/lib/api/contracts/knowledge/connectors.ts @@ -354,6 +354,7 @@ const searchSourceSummaryFields = { z.object({ credentialId: z.string().min(1).max(128), displayName: z.string(), + status: z.enum(['active', 'needs_reauth']).optional(), }) ) .max(SEARCH_SOURCE_CANDIDATE_PAGE_SIZE), diff --git a/apps/sim/lib/api/contracts/knowledge/personal-integrations.ts b/apps/sim/lib/api/contracts/knowledge/personal-integrations.ts new file mode 100644 index 00000000000..c8caf28fc87 --- /dev/null +++ b/apps/sim/lib/api/contracts/knowledge/personal-integrations.ts @@ -0,0 +1,89 @@ +import { z } from 'zod' +import { defineRouteContract } from '@/lib/api/contracts' +import { successResponseSchema } from '@/lib/api/contracts/knowledge/shared' +import { organizationIdSchema } from '@/lib/api/contracts/primitives' +import { searchConnectionTargetSchema } from '@/lib/knowledge/search/connection-target' + +export const personalSearchIntegrationSchema = z.object({ + name: z.string().max(200), + providerId: z.string().min(1).max(100), + connectorType: z.string().min(1).max(100), + connectorId: z.string().min(1).max(200), + knowledgeBaseId: z.string().min(1).max(200), + description: z.string().max(240), + accounts: z + .array( + z.object({ + credentialId: z.string().min(1).max(128), + displayName: z.string(), + status: z.enum(['connected', 'reconnect_needed']), + action: searchConnectionTargetSchema.nullable(), + }) + ) + .max(100), + connectionStatus: z.enum(['connected', 'reconnect_needed', 'not_connected', 'unavailable']), + indexingStatus: z.enum(['indexing', 'indexed', 'not_indexed', 'sync_failed', 'paused']), + searchableDocuments: z.number().int().nonnegative(), + action: searchConnectionTargetSchema.nullable(), +}) + +export const personalSearchIntegrationPageSchema = z.object({ + completedCredentialId: z.string().min(1).max(128).nullable(), + connections: z.array(personalSearchIntegrationSchema).max(50), + available: z + .array( + z.object({ + name: z.string().max(200), + description: z.string().max(240), + target: searchConnectionTargetSchema, + }) + ) + .max(200), + nextCursor: z.string().max(1024).nullable(), +}) +export type PersonalSearchIntegration = z.output +export type PersonalSearchIntegrationPage = z.output + +export const personalSearchIntegrationsQuerySchema = z.object({ + completionId: z.string().uuid().optional(), + organizationId: organizationIdSchema, + connectorType: z.string().trim().min(1).max(100).optional(), + connectorId: z.string().min(1).max(200).optional(), + cursor: z.string().min(1).max(1024).optional(), +}) +export type PersonalSearchIntegrationsQuery = z.input + +export const listPersonalSearchIntegrationsContract = defineRouteContract({ + method: 'GET', + path: '/api/knowledge/sim-search/personal-integrations', + query: personalSearchIntegrationsQuerySchema, + response: { mode: 'json', schema: successResponseSchema(personalSearchIntegrationPageSchema) }, +}) + +export const connectPersonalSearchIntegrationBodySchema = z.object({ + organizationId: organizationIdSchema, + target: searchConnectionTargetSchema, + oauthCompletionId: z.string().uuid(), + sourceConfig: z + .record(z.string().min(1).max(100), z.string().max(2000)) + .refine((config) => Object.keys(config).length <= 30, 'Too many source configuration fields') + .optional(), +}) +export type ConnectPersonalSearchIntegrationBody = z.input< + typeof connectPersonalSearchIntegrationBodySchema +> +export const connectPersonalSearchIntegrationContract = defineRouteContract({ + method: 'POST', + path: '/api/knowledge/sim-search/personal-integrations', + body: connectPersonalSearchIntegrationBodySchema, + response: { + mode: 'json', + schema: successResponseSchema( + z.object({ + url: z.string().url(), + connectorId: z.string().min(1).max(200), + knowledgeBaseId: z.string().min(1).max(200), + }) + ), + }, +}) diff --git a/apps/sim/lib/copilot/assistant/tool-policy.ts b/apps/sim/lib/copilot/assistant/tool-policy.ts index 4289ad9660d..2ee6fe3c3b3 100644 --- a/apps/sim/lib/copilot/assistant/tool-policy.ts +++ b/apps/sim/lib/copilot/assistant/tool-policy.ts @@ -3,6 +3,7 @@ import type { ToolMetadata } from '@/tools/metadata' export const ASSISTANT_TOOLS = new Set([ 'search_workspace', 'read_document', + 'list_integrations', 'search_integration_tools', 'call_integration_tool', 'oauth_get_auth_link', diff --git a/apps/sim/lib/copilot/generated/tool-catalog-v1.ts b/apps/sim/lib/copilot/generated/tool-catalog-v1.ts index b65bebaf7d1..614726dfe67 100644 --- a/apps/sim/lib/copilot/generated/tool-catalog-v1.ts +++ b/apps/sim/lib/copilot/generated/tool-catalog-v1.ts @@ -74,6 +74,7 @@ export interface ToolCatalogEntry { | 'knowledge' | 'list_deployment_versions' | 'list_integration_tools' + | 'list_integrations' | 'list_workspace_mcp_servers' | 'load_deployment' | 'load_integration_tool' @@ -212,6 +213,7 @@ export interface ToolCatalogEntry { | 'knowledge' | 'list_deployment_versions' | 'list_integration_tools' + | 'list_integrations' | 'list_workspace_mcp_servers' | 'load_deployment' | 'load_integration_tool' @@ -3541,6 +3543,21 @@ export const ListIntegrationTools: ToolCatalogEntry = { }, } +export const ListIntegrations: ToolCatalogEntry = { + id: 'list_integrations', + name: 'list_integrations', + route: 'sim', + mode: 'async', + parameters: { + additionalProperties: false, + properties: { + connectorType: { maxLength: 100, minLength: 1, type: 'string' }, + cursor: { maxLength: 1024, minLength: 1, type: 'string' }, + }, + type: 'object', + }, +} + export const ListWorkspaceMcpServers: ToolCatalogEntry = { id: 'list_workspace_mcp_servers', name: 'list_workspace_mcp_servers', @@ -7647,6 +7664,7 @@ export const TOOL_CATALOG: Record = { [Knowledge.id]: Knowledge, [ListDeploymentVersions.id]: ListDeploymentVersions, [ListIntegrationTools.id]: ListIntegrationTools, + [ListIntegrations.id]: ListIntegrations, [ListWorkspaceMcpServers.id]: ListWorkspaceMcpServers, [LoadDeployment.id]: LoadDeployment, [LoadIntegrationTool.id]: LoadIntegrationTool, diff --git a/apps/sim/lib/copilot/generated/tool-schemas-v1.ts b/apps/sim/lib/copilot/generated/tool-schemas-v1.ts index 388d341ae1f..e7fb2ccfda7 100644 --- a/apps/sim/lib/copilot/generated/tool-schemas-v1.ts +++ b/apps/sim/lib/copilot/generated/tool-schemas-v1.ts @@ -3474,6 +3474,25 @@ export const TOOL_RUNTIME_SCHEMAS: Record = { }, resultSchema: undefined, }, + list_integrations: { + parameters: { + additionalProperties: false, + properties: { + connectorType: { + maxLength: 100, + minLength: 1, + type: 'string', + }, + cursor: { + maxLength: 1024, + minLength: 1, + type: 'string', + }, + }, + type: 'object', + }, + resultSchema: undefined, + }, list_workspace_mcp_servers: { parameters: { type: 'object', diff --git a/apps/sim/lib/copilot/tool-executor/executor.ts b/apps/sim/lib/copilot/tool-executor/executor.ts index a8d407d9bad..bf024181716 100644 --- a/apps/sim/lib/copilot/tool-executor/executor.ts +++ b/apps/sim/lib/copilot/tool-executor/executor.ts @@ -48,11 +48,12 @@ export async function executeTool( (context.workspaceId || context.workflowId || context.requestMode !== 'assistant' || - !['search_workspace', 'read_document'].includes(toolId)) + !['search_workspace', 'read_document', 'list_integrations'].includes(toolId)) ) { return { success: false, - error: 'Organization Assistant can search and read connected documents.', + error: + 'Organization Assistant can search documents and inspect personal Search integrations.', } } if (context.requestMode === 'assistant' && !ASSISTANT_TOOLS.has(toolId)) { diff --git a/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.test.ts b/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.test.ts new file mode 100644 index 00000000000..000825ddc6b --- /dev/null +++ b/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.test.ts @@ -0,0 +1,73 @@ +/** @vitest-environment node */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const m = vi.hoisted(() => ({ read: vi.fn(), authorizeChat: vi.fn() })) +vi.mock('@/lib/copilot/chat/organization-chats', () => ({ + authorizeOrganizationChatDelegation: { execute: m.authorizeChat }, +})) +vi.mock('@/lib/knowledge/application/personal-search-integrations', () => ({ + listPersonalSearchIntegrations: { + get operation() { + return knowledgeOperations.listPersonalSearchIntegrations + }, + execute: m.read, + }, +})) + +import { listIntegrationsServerTool } from '@/lib/copilot/tools/server/knowledge/list-integrations' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' + +const context = { + userId: 'person', + organizationId: 'org', + chatId: 'private-chat', + toolCallId: 'call', + requestMode: 'assistant', + copilotToolExecution: true, +} +beforeEach(() => { + vi.clearAllMocks() + m.authorizeChat.mockResolvedValue(undefined) + m.read.mockResolvedValue({ connections: [], available: [], nextCursor: null }) +}) +describe('list_integrations', () => { + it('binds the read to the authorized current person and private organization chat', async () => { + expect( + await listIntegrationsServerTool.execute({ connectorType: 'gmail', cursor: 'page' }, context) + ).toMatchObject({ success: true }) + expect(m.read).toHaveBeenCalledWith({ + principal: expect.objectContaining({ + subjectUserId: 'person', + organizationId: 'org', + resourceScope: { chatId: 'private-chat' }, + }), + input: { organizationId: 'org', connectorType: 'gmail', cursor: 'page' }, + }) + expect(m.authorizeChat).toHaveBeenCalledOnce() + }) + it.each([ + { organizationId: 'forged' }, + { userId: 'another' }, + { connectorType: 'x'.repeat(101) }, + { cursor: 'x'.repeat(1025) }, + ])('rejects forged scope and unbounded arguments', async (args) => { + expect(await listIntegrationsServerTool.execute(args, context)).toMatchObject({ + success: false, + }) + expect(m.read).not.toHaveBeenCalled() + }) + it.each([ + { ...context, workspaceId: 'workspace' }, + { ...context, requestMode: 'agent' }, + { ...context, chatId: undefined }, + { ...context, copilotToolExecution: false }, + ])('rejects untrusted or non-organization contexts', async (invalid) => { + expect(await listIntegrationsServerTool.execute({}, invalid)).toMatchObject({ success: false }) + expect(m.read).not.toHaveBeenCalled() + }) + it('checks revoked chat access before returning any account data', async () => { + m.authorizeChat.mockRejectedValue(new Error('revoked')) + expect(await listIntegrationsServerTool.execute({}, context)).toMatchObject({ success: false }) + expect(m.read).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.ts b/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.ts new file mode 100644 index 00000000000..eae98d2763a --- /dev/null +++ b/apps/sim/lib/copilot/tools/server/knowledge/list-integrations.ts @@ -0,0 +1,52 @@ +import { createLogger } from '@sim/logger' +import { z } from 'zod' +import { + executeCopilotOrganizationKnowledgeUseCase, + messageForCopilotKnowledgeError, + requireCopilotKnowledgeScope, +} from '@/lib/copilot/application/execute-knowledge-use-case' +import type { BaseServerTool } from '@/lib/copilot/tools/server/base-tool' +import { listPersonalSearchIntegrations } from '@/lib/knowledge/application/personal-search-integrations' + +const logger = createLogger('ListSearchIntegrations') +const inputSchema = z + .object({ + connectorType: z.string().trim().min(1).max(100).optional(), + cursor: z.string().min(1).max(1024).optional(), + }) + .strict() + +export const listIntegrationsServerTool: BaseServerTool = { + name: 'list_integrations', + async execute(raw, context) { + try { + const scope = requireCopilotKnowledgeScope(context) + if (scope.kind !== 'organization') + throw new Error('Integration inventory requires organization Search') + const input = inputSchema.parse(raw) + const data = await executeCopilotOrganizationKnowledgeUseCase( + context, + listPersonalSearchIntegrations, + { + ...input, + organizationId: scope.organizationId, + } + ) + return { + success: true, + data, + message: + 'These are your current Search connections. Connected does not mean indexed. To offer a connection, emit the exact action or target inside a terminal tag, without a URL. Refresh this inventory after the user submits connection status. Follow nextCursor before claiming this list is complete.', + } + } catch (error) { + logger.error('Could not list personal Search integrations', { error }) + return { + success: false, + message: + error instanceof z.ZodError + ? 'Invalid integration inventory arguments' + : messageForCopilotKnowledgeError(error), + } + } + }, +} diff --git a/apps/sim/lib/copilot/tools/server/router.ts b/apps/sim/lib/copilot/tools/server/router.ts index 0c499a66e9e..e8506a24464 100644 --- a/apps/sim/lib/copilot/tools/server/router.ts +++ b/apps/sim/lib/copilot/tools/server/router.ts @@ -44,6 +44,7 @@ import { workspaceFileServerTool } from '@/lib/copilot/tools/server/files/worksp import { validateGeneratedToolPayload } from '@/lib/copilot/tools/server/generated-schema' import { generateImageServerTool } from '@/lib/copilot/tools/server/image/generate-image' import { knowledgeBaseServerTool } from '@/lib/copilot/tools/server/knowledge/knowledge-base' +import { listIntegrationsServerTool } from '@/lib/copilot/tools/server/knowledge/list-integrations' import { searchKnowledgeBaseServerTool } from '@/lib/copilot/tools/server/knowledge/search-knowledge-base' import { readDocumentServerTool, @@ -182,6 +183,7 @@ const baseServerToolRegistry: Record = { [getCredentialsServerTool.name]: getCredentialsServerTool, [knowledgeBaseServerTool.name]: knowledgeBaseServerTool, [searchKnowledgeBaseServerTool.name]: searchKnowledgeBaseServerTool, + [listIntegrationsServerTool.name]: listIntegrationsServerTool, [searchWorkspaceServerTool.name]: searchWorkspaceServerTool, [readDocumentServerTool.name]: readDocumentServerTool, [enrichmentRunServerTool.name]: enrichmentRunServerTool, diff --git a/apps/sim/lib/copilot/tools/tool-display.ts b/apps/sim/lib/copilot/tools/tool-display.ts index da2408eb90b..166840b8c22 100644 --- a/apps/sim/lib/copilot/tools/tool-display.ts +++ b/apps/sim/lib/copilot/tools/tool-display.ts @@ -565,6 +565,7 @@ const TOOL_TITLES: Record = { edit_workflow: 'Editing workflow', manage_knowledge_base: 'Managing knowledge base', search_knowledge_base: 'Searching knowledge base', + list_integrations: 'Checking your integrations', search_workspace: 'Searching documents', read_document: 'Reading document', open_resource: 'Opening resource', diff --git a/apps/sim/lib/credential-groups/oauth-intent.ts b/apps/sim/lib/credential-groups/oauth-intent.ts new file mode 100644 index 00000000000..5871ea6af69 --- /dev/null +++ b/apps/sim/lib/credential-groups/oauth-intent.ts @@ -0,0 +1,8 @@ +import { z } from 'zod' + +/** Keeps a connection card's create/reconnect intent intact through provider authorization. */ +export const credentialGroupConnectionIntentSchema = z.discriminatedUnion('kind', [ + z.object({ kind: z.literal('create') }).strict(), + z.object({ kind: z.literal('reconnect'), credentialId: z.string().min(1).max(128) }).strict(), +]) +export type CredentialGroupConnectionIntent = z.infer diff --git a/apps/sim/lib/credential-groups/oauth-state.test.ts b/apps/sim/lib/credential-groups/oauth-state.test.ts index 799f7157e65..2b9d7dc770e 100644 --- a/apps/sim/lib/credential-groups/oauth-state.test.ts +++ b/apps/sim/lib/credential-groups/oauth-state.test.ts @@ -70,6 +70,7 @@ describe('credential group OAuth state', () => { redirectUri: 'https://sim.ai/api/auth/oauth2/callback/google-email', codeVerifier: 'code-verifier', invitationToken: 'invitation-token', + connectionIntent: { kind: 'reconnect', credentialId: 'existing-account' }, completionRedirect: true, completionId: '550e8400-e29b-41d4-a716-446655440000', }) @@ -91,6 +92,7 @@ describe('credential group OAuth state', () => { optionId: 'option-1', codeVerifier: 'code-verifier', invitationToken: 'invitation-token', + connectionIntent: { kind: 'reconnect', credentialId: 'existing-account' }, completionRedirect: true, completionId: '550e8400-e29b-41d4-a716-446655440000', }) diff --git a/apps/sim/lib/credential-groups/oauth-state.ts b/apps/sim/lib/credential-groups/oauth-state.ts index 38f726ecfc6..edd588c1f37 100644 --- a/apps/sim/lib/credential-groups/oauth-state.ts +++ b/apps/sim/lib/credential-groups/oauth-state.ts @@ -5,6 +5,10 @@ import { getRedisClient } from '@/lib/core/config/redis' import { resourceScopeFields, resourceScopeFromOwner } from '@/lib/core/resource-scope' import { decryptSecret, encryptSecret } from '@/lib/core/security/encryption' import { assertCredentialGroupOAuthAttemptVersion } from '@/lib/credential-groups/oauth-attempt-version' +import { + type CredentialGroupConnectionIntent, + credentialGroupConnectionIntentSchema, +} from '@/lib/credential-groups/oauth-intent' import { type CredentialGroupProvider, isCredentialGroupProvider, @@ -38,6 +42,7 @@ interface StoredCredentialGroupOAuthAttempt { requiredScopes: string[] redirectUri: string completionRedirect?: boolean + connectionIntent?: CredentialGroupConnectionIntent completionId?: string returnTo?: 'search' | 'accounts' nonceHash: string @@ -62,6 +67,7 @@ export interface CredentialGroupOAuthAttempt { requiredScopes: string[] redirectUri: string completionRedirect?: boolean + connectionIntent?: CredentialGroupConnectionIntent completionId?: string returnTo?: 'search' | 'accounts' codeVerifier?: string @@ -83,6 +89,7 @@ interface CreateCredentialGroupOAuthAttemptParams { requiredScopes: string[] redirectUri: string completionRedirect?: boolean + connectionIntent?: CredentialGroupConnectionIntent completionId?: string returnTo?: 'search' | 'accounts' codeVerifier?: string @@ -136,6 +143,8 @@ function isStoredAttempt(value: unknown): value is StoredCredentialGroupOAuthAtt (candidate.completionRedirect === true && typeof candidate.completionId === 'string' && isValidUuid(candidate.completionId))) && + (candidate.connectionIntent === undefined || + credentialGroupConnectionIntentSchema.safeParse(candidate.connectionIntent).success) && (candidate.returnTo === undefined || candidate.returnTo === 'search' || candidate.returnTo === 'accounts') && @@ -157,6 +166,7 @@ export async function createCredentialGroupOAuthAttempt( ) { throw new Error('OAuth completion requires a valid correlation ID and completion redirect') } + if (params.connectionIntent) credentialGroupConnectionIntentSchema.parse(params.connectionIntent) const redis = requireRedis() const state = `${OAUTH_ATTEMPT_STATE_PREFIX}${generateId()}` const nonce = generateId() @@ -178,6 +188,7 @@ export async function createCredentialGroupOAuthAttempt( requiredScopes: params.requiredScopes, redirectUri: params.redirectUri, ...(params.completionRedirect ? { completionRedirect: true } : {}), + ...(params.connectionIntent ? { connectionIntent: params.connectionIntent } : {}), ...(params.completionId ? { completionId: params.completionId } : {}), ...(params.returnTo ? { returnTo: params.returnTo } : {}), nonceHash: sha256Hex(nonce), @@ -234,6 +245,7 @@ export async function consumeCredentialGroupOAuthAttempt( requiredScopes: parsed.requiredScopes, redirectUri: parsed.redirectUri, ...(parsed.completionRedirect ? { completionRedirect: true } : {}), + ...(parsed.connectionIntent ? { connectionIntent: parsed.connectionIntent } : {}), ...(parsed.completionId ? { completionId: parsed.completionId } : {}), ...(parsed.returnTo ? { returnTo: parsed.returnTo } : {}), ...(codeVerifier ? { codeVerifier: codeVerifier.decrypted } : {}), diff --git a/apps/sim/lib/credential-groups/oauth.test.ts b/apps/sim/lib/credential-groups/oauth.test.ts index b9879eba19f..0fae1303e67 100644 --- a/apps/sim/lib/credential-groups/oauth.test.ts +++ b/apps/sim/lib/credential-groups/oauth.test.ts @@ -277,6 +277,42 @@ describe('credential group OAuth persistence', () => { } ) + it.each([ + { kind: 'create' as const }, + { kind: 'reconnect' as const, credentialId: 'other-account' }, + ])( + 'rejects changing an active account through the wrong connection intent', + async (connectionIntent) => { + dbChainMockFns.limit.mockResolvedValueOnce([{ status: 'completed' }]) + queueTableRows(schemaMock.credentialGroup, [GROUP]) + queueTableRows(schemaMock.credential, [{ id: 'credential-1', revokedAt: null }]) + await expect( + completeCredentialGroupOAuth( + CONTEXT, + { + state: 'state', + userId: CONTEXT.credentialOwnerId, + provider: 'gmail', + nonceHash: 'nonce', + workspaceId: CONTEXT.workspaceId, + email: CONTEXT.email, + enrollmentId: CONTEXT.enrollmentId, + credentialGroupId: CONTEXT.credentialGroupId, + optionId: CONTEXT.option.id, + authorizationAppId: POLICY.authorizationAppId, + scopeVersion: POLICY.scopeVersion, + requiredScopes: POLICY.requiredScopes, + redirectUri: 'https://sim.test/callback', + invitationToken: 'token', + createdAt: Date.now(), + connectionIntent, + }, + 'code' + ) + ).rejects.toThrow('account changed') + expect(dbChainMockFns.update).not.toHaveBeenCalled() + } + ) it('preserves completed enrollment state when an account reconnects', async () => { dbChainMockFns.limit.mockResolvedValueOnce([{ status: 'completed' }]) queueTableRows(schemaMock.credentialGroup, [GROUP]) @@ -331,6 +367,61 @@ describe('credential group OAuth persistence', () => { }) }) + it('preserves the exact credential requested by Search reconnect', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ status: 'completed' }]) + queueTableRows(schemaMock.credentialGroup, [GROUP]) + queueTableRows(schemaMock.credential, [ + { + id: 'credential-1', + providerSubjectId: 'google-subject-1', + encryptedOauthTokenSet: null, + refreshTokenExpiresAt: null, + }, + ]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'credential-1' }]) + .mockResolvedValueOnce([{ id: CONTEXT.enrollmentId }]) + + const result = await completeCredentialGroupOAuth( + { ...CONTEXT, enrollmentStatus: 'completed' }, + { + connectionIntent: { kind: 'reconnect', credentialId: 'credential-1' }, + state: 'state-1', + userId: 'person-1', + provider: 'gmail', + nonceHash: 'nonce-hash', + workspaceId: CONTEXT.workspaceId, + email: CONTEXT.email, + enrollmentId: CONTEXT.enrollmentId, + credentialGroupId: CONTEXT.credentialGroupId, + optionId: CONTEXT.option.id, + authorizationAppId: POLICY.authorizationAppId, + scopeVersion: POLICY.scopeVersion, + requiredScopes: POLICY.requiredScopes, + redirectUri: 'https://sim.ai/api/auth/oauth2/callback/google-email', + codeVerifier: 'verifier', + invitationToken: 'invitation-token', + createdAt: Date.now(), + }, + 'authorization-code' + ) + + const enrollmentUpdate = dbChainMockFns.set.mock.calls[1]?.[0] + expect(enrollmentUpdate).toEqual( + expect.objectContaining({ status: 'completed', updatedAt: expect.any(Date) }) + ) + expect(enrollmentUpdate).not.toHaveProperty('completedAt') + expect(result).toEqual({ + created: false, + credentialId: 'credential-1', + credentialGroupOptionId: 'option-1', + provider: 'gmail', + providerId: 'google-email', + displayName: 'person@example.com', + enrollmentStatus: 'completed', + }) + }) + it('rejects an exchanged grant when the group policy changed before persistence', async () => { const nextPolicy = { ...POLICY, @@ -420,4 +511,59 @@ describe('credential group OAuth persistence', () => { expect(eq).toHaveBeenCalledWith(schemaMock.credentialGroupEnrollment.email, CONTEXT.email) expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) + it('restores a personally disconnected grant without replacing any active account', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ status: 'completed' }]) + queueTableRows(schemaMock.credentialGroup, [GROUP]) + queueTableRows(schemaMock.credential, [ + { + id: 'credential-1', + revokedAt: new Date(), + providerSubjectId: 'google-subject-1', + encryptedOauthTokenSet: null, + refreshTokenExpiresAt: null, + }, + ]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'credential-1' }]) + .mockResolvedValueOnce([{ id: CONTEXT.enrollmentId }]) + + const result = await completeCredentialGroupOAuth( + { ...CONTEXT, enrollmentStatus: 'completed' }, + { + connectionIntent: { kind: 'create' }, + state: 'state-1', + userId: 'person-1', + provider: 'gmail', + nonceHash: 'nonce-hash', + workspaceId: CONTEXT.workspaceId, + email: CONTEXT.email, + enrollmentId: CONTEXT.enrollmentId, + credentialGroupId: CONTEXT.credentialGroupId, + optionId: CONTEXT.option.id, + authorizationAppId: POLICY.authorizationAppId, + scopeVersion: POLICY.scopeVersion, + requiredScopes: POLICY.requiredScopes, + redirectUri: 'https://sim.ai/api/auth/oauth2/callback/google-email', + codeVerifier: 'verifier', + invitationToken: 'invitation-token', + createdAt: Date.now(), + }, + 'authorization-code' + ) + + const enrollmentUpdate = dbChainMockFns.set.mock.calls[1]?.[0] + expect(enrollmentUpdate).toEqual( + expect.objectContaining({ status: 'completed', updatedAt: expect.any(Date) }) + ) + expect(enrollmentUpdate).not.toHaveProperty('completedAt') + expect(result).toEqual({ + created: false, + credentialId: 'credential-1', + credentialGroupOptionId: 'option-1', + provider: 'gmail', + providerId: 'google-email', + displayName: 'person@example.com', + enrollmentStatus: 'completed', + }) + }) }) diff --git a/apps/sim/lib/credential-groups/oauth.ts b/apps/sim/lib/credential-groups/oauth.ts index a6c2a0858b1..b905f373604 100644 --- a/apps/sim/lib/credential-groups/oauth.ts +++ b/apps/sim/lib/credential-groups/oauth.ts @@ -16,6 +16,7 @@ import { type CredentialGroupOAuthContext, lockCredentialGroupEnrollmentLifecycle, } from '@/lib/credential-groups/enrollments' +import type { CredentialGroupConnectionIntent } from '@/lib/credential-groups/oauth-intent' import { type CredentialGroupOAuthAttempt, createCredentialGroupOAuthAttempt, @@ -35,6 +36,7 @@ import { getCredentialGroupProviderService, isCredentialGroupProvider, } from '@/lib/credential-groups/providers' +import { recordSearchConnectionCompletion } from '@/lib/credential-groups/search-connection-completion' import { decryptManagedOAuthTokenSet, encryptManagedOAuthTokenSet, @@ -112,6 +114,7 @@ export async function startCredentialGroupOAuth( invitationToken: string, options: { completionRedirect?: boolean + connectionIntent?: CredentialGroupConnectionIntent completionId?: string returnTo?: 'search' | 'accounts' } = {} @@ -135,6 +138,7 @@ export async function startCredentialGroupOAuth( codeVerifier: prepared.codeVerifier, completionRedirect: options.completionRedirect, completionId: options.completionId, + connectionIntent: options.connectionIntent, returnTo: options.returnTo, invitationToken, }) @@ -146,7 +150,8 @@ async function persistGrant( adapter: CredentialGroupProviderAdapter, policy: CredentialGroupProviderPolicy, grant: VerifiedCredentialGroupGrant, - invitationTokenHash: string + invitationTokenHash: string, + connectionIntent?: CredentialGroupConnectionIntent ): Promise { if (grant.providerId !== policy.providerId) { throw new CredentialGroupOAuthError('Provider returned a credential for another app.', 502) @@ -217,6 +222,7 @@ async function persistGrant( const [existing] = await tx .select({ id: credential.id, + revokedAt: credential.revokedAt, providerSubjectId: credential.providerSubjectId, encryptedOauthTokenSet: credential.encryptedOauthTokenSet, refreshTokenExpiresAt: credential.refreshTokenExpiresAt, @@ -231,6 +237,16 @@ async function persistGrant( ) .limit(1) + if ( + (connectionIntent?.kind === 'create' && existing && !existing.revokedAt) || + (connectionIntent?.kind === 'reconnect' && existing?.id !== connectionIntent.credentialId) + ) { + throw new CredentialGroupOAuthError( + 'The account changed during authorization. Refresh your connections and try again.', + 409 + ) + } + let refreshToken = grant.refreshToken if ( !refreshToken && @@ -373,5 +389,21 @@ export async function completeCredentialGroupOAuth( const adapter = getOptionAdapter(context) const policy = await assertCurrentPolicy(context, adapter, attempt) const grant = await adapter.exchangeAndVerify({ context, attempt, code, policy }) - return persistGrant(context, adapter, policy, grant, sha256Hex(attempt.invitationToken)) + const completion = await persistGrant( + context, + adapter, + policy, + grant, + sha256Hex(attempt.invitationToken), + attempt.connectionIntent + ) + if (attempt.connectionIntent && attempt.completionId && attempt.organizationId) { + await recordSearchConnectionCompletion({ + organizationId: attempt.organizationId, + userId: attempt.userId, + completionId: attempt.completionId, + credentialId: completion.credentialId, + }) + } + return completion } diff --git a/apps/sim/lib/credential-groups/search-connection-completion.test.ts b/apps/sim/lib/credential-groups/search-connection-completion.test.ts new file mode 100644 index 00000000000..b46b633b2f3 --- /dev/null +++ b/apps/sim/lib/credential-groups/search-connection-completion.test.ts @@ -0,0 +1,54 @@ +/** @vitest-environment node */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const m = vi.hoisted(() => { + const values = new Map() + return { + values, + redis: { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn(async (key: string, value: string) => { + if (values.has(key)) return null + values.set(key, value) + return 'OK' + }), + }, + } +}) +vi.mock('@/lib/core/config/redis', () => ({ getRedisClient: () => m.redis })) + +import { + readSearchConnectionCompletion, + recordSearchConnectionCompletion, +} from '@/lib/credential-groups/search-connection-completion' + +const scope = { userId: 'person', organizationId: 'org', completionId: 'attempt' } +beforeEach(() => { + vi.clearAllMocks() + m.values.clear() +}) +describe('Search OAuth completion receipts', () => { + it('isolates people, organizations and concurrent attempts', async () => { + await recordSearchConnectionCompletion({ ...scope, credentialId: 'mine' }) + expect(await readSearchConnectionCompletion(scope)).toBe('mine') + expect(await readSearchConnectionCompletion({ ...scope, userId: 'someone-else' })).toBeNull() + expect(await readSearchConnectionCompletion({ ...scope, organizationId: 'another' })).toBeNull() + expect( + await readSearchConnectionCompletion({ ...scope, completionId: 'other-attempt' }) + ).toBeNull() + expect(m.redis.set).toHaveBeenCalledWith( + expect.any(String), + JSON.stringify({ credentialId: 'mine' }), + 'EX', + 86_400, + 'NX' + ) + }) + it('does not allow a completion to be overwritten', async () => { + await recordSearchConnectionCompletion({ ...scope, credentialId: 'mine' }) + await expect( + recordSearchConnectionCompletion({ ...scope, credentialId: 'changed' }) + ).rejects.toThrow('already recorded') + expect(await readSearchConnectionCompletion(scope)).toBe('mine') + }) +}) diff --git a/apps/sim/lib/credential-groups/search-connection-completion.ts b/apps/sim/lib/credential-groups/search-connection-completion.ts new file mode 100644 index 00000000000..7812ec247d3 --- /dev/null +++ b/apps/sim/lib/credential-groups/search-connection-completion.ts @@ -0,0 +1,40 @@ +import { sha256Hex } from '@sim/security/hash' +import { z } from 'zod' +import { getRedisClient } from '@/lib/core/config/redis' + +interface SearchConnectionCompletionScope { + organizationId: string + userId: string + completionId: string +} +const receiptSchema = z.object({ credentialId: z.string().min(1).max(128) }).strict() + +function completionKey(scope: SearchConnectionCompletionScope) { + return `search:connection-completion:${sha256Hex(JSON.stringify([scope.organizationId, scope.userId, scope.completionId]))}` +} + +/** A short-lived receipt proves which OAuth attempt committed, even after a page reload. */ +export async function recordSearchConnectionCompletion( + input: SearchConnectionCompletionScope & { credentialId: string } +) { + const redis = getRedisClient() + if (!redis) throw new Error('Search connection completion requires Redis') + const stored = await redis.set( + completionKey(input), + JSON.stringify(receiptSchema.parse({ credentialId: input.credentialId })), + 'EX', + 86_400, + 'NX' + ) + if (stored !== 'OK') throw new Error('Search connection completion was already recorded') +} + +/** Called only inside the authorized personal inventory read; scope never comes from a tag. */ +export async function readSearchConnectionCompletion( + scope: SearchConnectionCompletionScope +): Promise { + const redis = getRedisClient() + if (!redis) throw new Error('Search connection completion requires Redis') + const value = await redis.get(completionKey(scope)) + return value ? receiptSchema.parse(JSON.parse(value)).credentialId : null +} diff --git a/apps/sim/lib/knowledge/application/connect-personal-search-integration.ts b/apps/sim/lib/knowledge/application/connect-personal-search-integration.ts new file mode 100644 index 00000000000..3da954fd42b --- /dev/null +++ b/apps/sim/lib/knowledge/application/connect-personal-search-integration.ts @@ -0,0 +1,37 @@ +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { resolveKnowledgeOrganizationContext } from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { resolvePersonalSearchConnection } from '@/lib/knowledge/application/personal-search-integrations' +import { connectSimSearchConnector } from '@/lib/knowledge/application/sim-search' +import type { SearchConnectionTarget } from '@/lib/knowledge/search/connection-target' + +interface ConnectPersonalSearchIntegrationInput { + organizationId: string + target: SearchConnectionTarget + oauthCompletionId: string + sourceConfig?: Record +} + +/** A user click validates the requested control before starting the ordinary source enrollment. */ +export const connectPersonalSearchIntegration = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.connectPersonalSearchIntegration, + resolveContext: ({ input }: { input: ConnectPersonalSearchIntegrationInput }) => + resolveKnowledgeOrganizationContext(input), + async execute({ principal, input, request }) { + const { target } = await resolvePersonalSearchConnection.execute({ principal, input }) + return connectSimSearchConnector.execute({ + principal, + request, + input: { + organizationId: input.organizationId, + connectorType: target.connectorType, + connectorId: target.connectorId, + sourceConfig: input.sourceConfig, + oauthCompletionId: input.oauthCompletionId, + connectionIntent: target.credentialId + ? { kind: 'reconnect', credentialId: target.credentialId } + : { kind: 'create' }, + }, + }) + }, +}) diff --git a/apps/sim/lib/knowledge/application/connector-access.ts b/apps/sim/lib/knowledge/application/connector-access.ts index aded446945e..b3a5eb19ec4 100644 --- a/apps/sim/lib/knowledge/application/connector-access.ts +++ b/apps/sim/lib/knowledge/application/connector-access.ts @@ -12,6 +12,7 @@ import { generateRequestId } from '@/lib/core/utils/request' import { loadScopedAccountsCredentialListContext } from '@/lib/credential-groups/credentials' import { getCredentialGroupOAuthContextForEnrollment } from '@/lib/credential-groups/enrollments' import { startCredentialGroupOAuth } from '@/lib/credential-groups/oauth' +import type { CredentialGroupConnectionIntent } from '@/lib/credential-groups/oauth-intent' import { createViewerCredentialGroupEnrollment } from '@/lib/credential-groups/self-enrollment' import { requireKnowledgeMemberAccessAvailable, @@ -58,6 +59,7 @@ export interface StartKnowledgeConnectorMemberEnrollmentInput { assertedWorkspaceId?: string assertedOrganizationId?: string /** Opens provider OAuth directly and correlates its completion with the initiating tab. */ + connectionIntent?: CredentialGroupConnectionIntent oauthCompletionId?: string } @@ -123,6 +125,7 @@ export const startKnowledgeConnectorMemberEnrollment = defineAuthorizedKnowledge completionRedirect: true, returnTo: 'search', completionId: input.oauthCompletionId, + connectionIntent: input.connectionIntent, }) } if (!context.knowledgeBase.isSearchIndex) return invitationLink diff --git a/apps/sim/lib/knowledge/application/operations.test.ts b/apps/sim/lib/knowledge/application/operations.test.ts index 246bcd4e307..1dc78ba611f 100644 --- a/apps/sim/lib/knowledge/application/operations.test.ts +++ b/apps/sim/lib/knowledge/application/operations.test.ts @@ -71,6 +71,8 @@ describe('knowledge operation registry', () => { 'knowledge.connectors.create', 'knowledge.connectors.update', 'knowledge.connectors.access.update', + 'knowledge.search.personal-integrations.connect', + 'knowledge.search.personal-integrations.list', 'knowledge.search.sources.list', 'knowledge.search.sources.overview', 'knowledge.search.sources.progress', diff --git a/apps/sim/lib/knowledge/application/operations.ts b/apps/sim/lib/knowledge/application/operations.ts index 0e7ae360969..e9a7e1cb9a7 100644 --- a/apps/sim/lib/knowledge/application/operations.ts +++ b/apps/sim/lib/knowledge/application/operations.ts @@ -692,13 +692,33 @@ export const knowledgeOperations = { principalKinds: ['session'], }) ), + connectPersonalSearchIntegration: defineKnowledgeOperation( + defineWorkspaceOperation({ + id: 'knowledge.search.personal-integrations.connect', + minimumRole: 'read', + workspaceApiKey: 'deny', + capability: 'knowledge.use', + principalKinds: ['session'], + }) + ), + listPersonalSearchIntegrations: defineKnowledgeOperation( + defineWorkspaceOperation({ + id: 'knowledge.search.personal-integrations.list', + minimumRole: 'read', + workspaceApiKey: 'deny', + capability: 'knowledge.use', + principalKinds: ['session', 'delegated'], + delegatedServices: ['copilot'], + }) + ), listSearchSources: defineKnowledgeOperation( defineWorkspaceOperation({ id: 'knowledge.search.sources.list', minimumRole: 'read', workspaceApiKey: 'deny', capability: 'knowledge.use', - principalKinds: ['session'], + principalKinds: ['session', 'delegated'], + delegatedServices: ['copilot'], }) ), readSearchSourceOverview: defineKnowledgeOperation( @@ -707,7 +727,8 @@ export const knowledgeOperations = { minimumRole: 'read', workspaceApiKey: 'deny', capability: 'knowledge.use', - principalKinds: ['session'], + principalKinds: ['session', 'delegated'], + delegatedServices: ['copilot'], }) ), readSearchSourceProgress: defineKnowledgeOperation( diff --git a/apps/sim/lib/knowledge/application/personal-search-integrations.test.ts b/apps/sim/lib/knowledge/application/personal-search-integrations.test.ts new file mode 100644 index 00000000000..4f2bf73e86b --- /dev/null +++ b/apps/sim/lib/knowledge/application/personal-search-integrations.test.ts @@ -0,0 +1,189 @@ +/** @vitest-environment node */ +import { user } from '@sim/db/schema' +import { queueTableRows, resetDbChainMock } from '@sim/testing' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const m = vi.hoisted(() => ({ + authorize: vi.fn(), + sources: vi.fn(), + overview: vi.fn(), + approvals: vi.fn(), + availability: vi.fn(), +})) +vi.mock('@/lib/core/application/organization-authorization', () => ({ + authorizeOrganizationOperation: m.authorize, +})) +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveKnowledgeOrganizationContext: async ({ organizationId }: { organizationId: string }) => ({ + organizationId, + workspaceId: undefined, + }), +})) +vi.mock('@/lib/knowledge/application/search-sources', () => ({ + listSearchSources: { execute: m.sources }, +})) +vi.mock('@/lib/knowledge/application/search-source-overview', () => ({ + readSearchSourceOverview: { execute: m.overview }, +})) +vi.mock('@/lib/knowledge/search/integration-policy', () => ({ + listOrganizationSearchApprovals: m.approvals, +})) +vi.mock('@/lib/knowledge/access/availability', () => ({ + resolveKnowledgeAccessAvailability: async () => ({ memberScoped: true, sourceMirrored: true }), +})) +vi.mock('@/lib/integrations/availability.server', () => ({ + getIntegrationAvailability: () => [], + isOAuthServiceDeploymentAvailable: () => true, +})) +vi.mock('@/lib/sim-search/connectors', () => ({ + SEARCH_CONNECTORS: ['gmail', 'slack', 'notion'].map((type) => ({ + type, + providerId: type, + meta: { name: type }, + setupFields: [], + })), + getConnectorAccessAvailability: m.availability, +})) + +import { personalSearchIntegrationPageSchema } from '@/lib/api/contracts/knowledge/personal-integrations' +import { + listPersonalSearchIntegrations, + resolvePersonalSearchConnection, +} from '@/lib/knowledge/application/personal-search-integrations' + +const principal = { kind: 'session', userId: 'person', sessionId: 'session' } as const +const input = { organizationId: 'org' } +const target = { + type: 'link', + provider: 'gmail', + connectorType: 'gmail', + connectorId: 'source', +} as const +const source = { + connectorType: 'gmail', + connectorId: 'source', + knowledgeBaseId: 'kb', + sourceDescription: '', + enabled: true, + availability: 'available', + viewerEmailVerified: true, + connectionRequired: true, + viewerMembership: 'connected', + viewerAccounts: [{ credentialId: 'mine', displayName: 'My mail', status: 'active' }], + isSyncing: false, + hasSyncError: false, + viewerFailedDocumentCount: 0, + viewerDocumentCount: 0, +} +beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + queueTableRows(user, [{ emailVerified: true }]) + m.authorize.mockResolvedValue(undefined) + m.sources.mockResolvedValue({ sources: [source], nextCursor: null }) + m.overview.mockResolvedValue({ providers: [{ connectorType: 'gmail' }] }) + m.approvals.mockResolvedValue( + new Map([ + ['gmail', true], + ['slack', true], + ]) + ) + m.availability.mockReturnValue({ members: true }) +}) +describe('personal Search inventory', () => { + it.each([ + [{}, 'connected', 'not_indexed'], + [{ isSyncing: true }, 'connected', 'indexing'], + [{ viewerDocumentCount: 3 }, 'connected', 'indexed'], + [{ hasSyncError: true }, 'connected', 'sync_failed'], + [ + { + viewerAccounts: [{ credentialId: 'mine', displayName: 'My mail', status: 'needs_reauth' }], + }, + 'reconnect_needed', + 'not_indexed', + ], + ])( + 'separates account state from index state %#', + async (changes, connectionStatus, indexingStatus) => { + m.sources.mockResolvedValue({ sources: [{ ...source, ...changes }], nextCursor: null }) + const result = await listPersonalSearchIntegrations.execute({ principal, input }) + expect(result.connections[0]).toMatchObject({ connectionStatus, indexingStatus }) + expect(personalSearchIntegrationPageSchema.safeParse(result).success).toBe(true) + expect(m.sources).toHaveBeenCalledWith({ principal, input }) + expect(JSON.stringify(result)).not.toMatch(/accessToken|authorizationUrl|other-person/) + } + ) + it('offers approved ready providers with no index, excluding app setup and unapproved providers', async () => { + m.sources.mockResolvedValue({ sources: [], nextCursor: null }) + m.overview.mockResolvedValue({ providers: [] }) + const result = await listPersonalSearchIntegrations.execute({ principal, input }) + expect(result.available.map((entry) => entry.target.connectorType)).toEqual(['gmail']) + expect(result.connections).toEqual([]) + }) + it('omits every connection action when the person has an unverified email', async () => { + resetDbChainMock() + queueTableRows(user, [{ emailVerified: false }]) + m.sources.mockResolvedValue({ sources: [], nextCursor: null }) + m.overview.mockResolvedValue({ providers: [] }) + expect((await listPersonalSearchIntegrations.execute({ principal, input })).available).toEqual( + [] + ) + }) + it('forwards the bounded cursor and provider filter without exposing other people’s accounts', async () => { + m.sources.mockResolvedValue({ + sources: [{ ...source, viewerAccounts: [], viewerMembership: 'not_enrolled' }], + nextCursor: 'next', + }) + const filtered = { ...input, connectorType: 'gmail', cursor: 'page' } + const result = await listPersonalSearchIntegrations.execute({ principal, input: filtered }) + expect(result.connections).toEqual([]) + expect(result.nextCursor).toBe('next') + expect(result.available).toEqual([{ name: 'gmail', description: '', target }]) + expect(m.sources).toHaveBeenCalledWith({ principal, input: filtered }) + }) + it('rechecks authorization before every read and returns nothing after membership is revoked', async () => { + m.authorize.mockRejectedValue(new Error('Membership revoked')) + await expect(listPersonalSearchIntegrations.execute({ principal, input })).rejects.toThrow( + 'Membership revoked' + ) + expect(m.sources).not.toHaveBeenCalled() + }) + it.each([ + { ...target, credentialId: 'another-person' }, + { ...target, connectorId: 'other-source' }, + { ...target, provider: 'notion' }, + ])('rejects a forged or stale reconnect target', async (forged) => { + m.sources.mockResolvedValue({ + sources: [ + { + ...source, + viewerAccounts: [ + { credentialId: 'mine', displayName: 'My mail', status: 'needs_reauth' }, + ], + }, + ], + nextCursor: null, + }) + await expect( + resolvePersonalSearchConnection.execute({ principal, input: { ...input, target: forged } }) + ).rejects.toThrow('no longer available') + }) + it('returns the precise owned reconnect target', async () => { + m.sources.mockResolvedValue({ + sources: [ + { + ...source, + viewerAccounts: [ + { credentialId: 'mine', displayName: 'My mail', status: 'needs_reauth' }, + ], + }, + ], + nextCursor: null, + }) + const selected = { ...target, credentialId: 'mine' } + await expect( + resolvePersonalSearchConnection.execute({ principal, input: { ...input, target: selected } }) + ).resolves.toEqual({ name: 'gmail', target: selected }) + }) +}) diff --git a/apps/sim/lib/knowledge/application/personal-search-integrations.ts b/apps/sim/lib/knowledge/application/personal-search-integrations.ts new file mode 100644 index 00000000000..735f5acd97f --- /dev/null +++ b/apps/sim/lib/knowledge/application/personal-search-integrations.ts @@ -0,0 +1,214 @@ +import { requirePrincipalSubjectUserId } from '@sim/auth/principal' +import { db } from '@sim/db' +import { user } from '@sim/db/schema' +import { eq } from 'drizzle-orm' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { readSearchConnectionCompletion } from '@/lib/credential-groups/search-connection-completion' +import { + getIntegrationAvailability, + isOAuthServiceDeploymentAvailable, +} from '@/lib/integrations/availability.server' +import { resolveKnowledgeAccessAvailability } from '@/lib/knowledge/access/availability' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { resolveKnowledgeOrganizationContext } from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { readSearchSourceOverview } from '@/lib/knowledge/application/search-source-overview' +import { listSearchSources } from '@/lib/knowledge/application/search-sources' +import type { SearchConnectionTarget } from '@/lib/knowledge/search/connection-target' +import { listOrganizationSearchApprovals } from '@/lib/knowledge/search/integration-policy' +import { getConnectorAccessAvailability, SEARCH_CONNECTORS } from '@/lib/sim-search/connectors' + +export interface ListPersonalSearchIntegrationsInput { + organizationId: string + connectorType?: string + connectorId?: string + cursor?: string + completionId?: string +} + +/** Current personal connections and eligible setup controls, shared by chat and Integrations. */ +export const listPersonalSearchIntegrations = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.listPersonalSearchIntegrations, + resolveContext: ({ input }: { input: ListPersonalSearchIntegrationsInput }) => + resolveKnowledgeOrganizationContext(input), + async execute({ principal, input, context }) { + const userId = requirePrincipalSubjectUserId(principal) + const [viewer] = await db + .select({ emailVerified: user.emailVerified }) + .from(user) + .where(eq(user.id, userId)) + .limit(1) + if (!viewer) throw new OrchestrationError('forbidden', 'The current person is unavailable') + const [page, overview, approvals, access] = await Promise.all([ + listSearchSources.execute({ principal, input }), + readSearchSourceOverview.execute({ + principal, + input: { organizationId: context.organizationId }, + }), + listOrganizationSearchApprovals(context.organizationId), + resolveKnowledgeAccessAvailability(context), + ]) + const deployment = new Map( + getIntegrationAvailability().map((entry) => [entry.type.toLowerCase(), entry]) + ) + const oauth = new Map( + SEARCH_CONNECTORS.map((entry) => [ + entry.providerId, + isOAuthServiceDeploymentAvailable(entry.providerId), + ]) + ) + const configured = new Set(overview.providers.map((entry) => entry.connectorType)) + const eligible = (connectorType: string) => { + const connector = SEARCH_CONNECTORS.find((entry) => entry.type === connectorType) + return Boolean( + viewer.emailVerified && + connector && + approvals.get(connectorType) && + getConnectorAccessAvailability(connector.meta, deployment, { + memberAccessAvailable: access.memberScoped, + mirroredAccessAvailable: access.sourceMirrored, + oauthServiceAvailability: oauth, + isIntegrationAvailabilityReady: true, + }).members + ) + } + const projected = page.sources.flatMap((source) => { + const connector = SEARCH_CONNECTORS.find((entry) => entry.type === source.connectorType) + if (!connector) return [] + const target: SearchConnectionTarget = { + type: 'link', + provider: connector.providerId, + connectorType: source.connectorType, + connectorId: source.connectorId, + } + const canConnect = + eligible(source.connectorType) && + source.enabled && + source.availability === 'available' && + source.viewerEmailVerified && + source.connectionRequired && + source.viewerMembership !== null && + !['revoked', 'unverified_email'].includes(source.viewerMembership) + const accounts = source.viewerAccounts.map((account) => { + if (!account.status) throw new Error('Personal Search account status is missing') + return { + credentialId: account.credentialId, + displayName: account.displayName, + status: + account.status === 'active' ? ('connected' as const) : ('reconnect_needed' as const), + action: + canConnect && account.status === 'needs_reauth' + ? { ...target, credentialId: account.credentialId } + : null, + } + }) + return [ + { + name: connector.meta.name, + providerId: connector.providerId, + connectorType: connector.type, + connectorId: source.connectorId, + knowledgeBaseId: source.knowledgeBaseId, + description: source.sourceDescription, + accounts, + connectionStatus: accounts.some((account) => account.status === 'reconnect_needed') + ? ('reconnect_needed' as const) + : accounts.length + ? ('connected' as const) + : canConnect + ? ('not_connected' as const) + : ('unavailable' as const), + indexingStatus: + !source.enabled || source.availability !== 'available' || source.approved === false + ? ('paused' as const) + : source.isSyncing + ? ('indexing' as const) + : source.hasSyncError || source.viewerFailedDocumentCount > 0 + ? ('sync_failed' as const) + : source.viewerDocumentCount > 0 + ? ('indexed' as const) + : ('not_indexed' as const), + searchableDocuments: source.viewerDocumentCount, + action: canConnect && !accounts.length ? target : null, + }, + ] + }) + const available: Array<{ name: string; description: string; target: SearchConnectionTarget }> = + [ + ...projected.flatMap((entry) => + entry.action + ? [{ name: entry.name, description: entry.description, target: entry.action }] + : [] + ), + ...SEARCH_CONNECTORS.filter( + (connector) => + !input.connectorId && + (!input.connectorType || connector.type === input.connectorType) && + connector.type !== 'slack' && + (!configured.has(connector.type) || connector.setupFields.length > 0) && + eligible(connector.type) + ).map((connector) => ({ + name: connector.meta.name, + description: '', + target: { + type: 'link' as const, + provider: connector.providerId, + connectorType: connector.type, + }, + })), + ] + return { + completedCredentialId: input.completionId + ? await readSearchConnectionCompletion({ + organizationId: context.organizationId, + userId, + completionId: input.completionId, + }) + : null, + connections: projected.filter((entry) => entry.accounts.length > 0), + available, + nextCursor: page.nextCursor, + } + }, +}) + +/** Revalidates a model- or URL-supplied target against current personal eligibility. */ +export const resolvePersonalSearchConnection = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.listPersonalSearchIntegrations, + resolveContext: ({ + input, + }: { + input: { organizationId: string; target: SearchConnectionTarget } + }) => resolveKnowledgeOrganizationContext(input), + async execute({ principal, input }) { + const page = await listPersonalSearchIntegrations.execute({ + principal, + input: { + organizationId: input.organizationId, + connectorType: input.target.connectorType, + connectorId: input.target.connectorId, + }, + }) + const targets = [ + ...page.available.map((entry) => ({ name: entry.name, target: entry.target })), + ...page.connections.flatMap((entry) => + entry.accounts.flatMap((account) => + account.action ? [{ name: entry.name, target: account.action }] : [] + ) + ), + ] + const selected = targets.find( + ({ target }) => + target.provider === input.target.provider && + target.connectorType === input.target.connectorType && + target.connectorId === input.target.connectorId && + target.credentialId === input.target.credentialId + ) + if (!selected) + throw new OrchestrationError( + 'validation', + 'This connection is no longer available. Refresh your integrations and try again.' + ) + return selected + }, +}) diff --git a/apps/sim/lib/knowledge/application/search-sources.ts b/apps/sim/lib/knowledge/application/search-sources.ts index b41e7a80270..8f33d05cfe1 100644 --- a/apps/sim/lib/knowledge/application/search-sources.ts +++ b/apps/sim/lib/knowledge/application/search-sources.ts @@ -1,3 +1,4 @@ +import { requirePrincipalSubjectUserId } from '@sim/auth/principal' import { db } from '@sim/db' import { document, embedding, knowledgeBase, knowledgeConnector, user } from '@sim/db/schema' import { and, desc, eq, exists, inArray, isNull, lt, or, sql } from 'drizzle-orm' @@ -27,6 +28,7 @@ import { getConnectorMeta } from '@/connectors/registry' export interface ListSearchSourcesInput extends ResourceOwner { cursor?: string + connectorId?: string connectorType?: string search?: string mine?: boolean @@ -38,14 +40,16 @@ export const listSearchSources = defineAuthorizedKnowledgeUseCase({ resolveContext: ({ input }: { input: ListSearchSourcesInput }) => resolveKnowledgeOwnerContext(input), async execute({ principal, input, context }) { + const userId = requirePrincipalSubjectUserId(principal) const search = input.search?.trim().toLowerCase() ?? '' const connectorType = input.connectorType?.trim() const cursorScope = cursorScopeKey(cursorRoute(listSearchSourcesContract), { workspaceId: context.workspaceId, organizationId: context.organizationId, - userId: principal.userId, + userId: userId, search, connectorType: connectorType ?? '', + connectorId: input.connectorId ?? '', mine: input.mine === true, order: 'newest', }) @@ -91,6 +95,7 @@ export const listSearchSources = defineAuthorizedKnowledgeUseCase({ isNull(knowledgeConnector.archivedAt), isNull(knowledgeConnector.deletedAt), connectorType ? eq(knowledgeConnector.connectorType, connectorType) : undefined, + input.connectorId ? eq(knowledgeConnector.id, input.connectorId) : undefined, cursor ? or( sql`${knowledgeConnector.createdAt} < ${cursor.createdAt}::timestamp`, @@ -110,7 +115,7 @@ export const listSearchSources = defineAuthorizedKnowledgeUseCase({ const [availability, memberships, viewers, approvals, accounts] = await Promise.all([ resolveKnowledgeAccessAvailability(context), resolveViewerConnectorMemberships({ - userId: principal.userId, + userId: userId, workspaceId: context.workspaceId, organizationId: context.organizationId, connectors: scanned, @@ -118,13 +123,13 @@ export const listSearchSources = defineAuthorizedKnowledgeUseCase({ db .select({ emailVerified: user.emailVerified }) .from(user) - .where(eq(user.id, principal.userId)) + .where(eq(user.id, userId)) .limit(1), context.organizationId ? listOrganizationSearchApprovals(context.organizationId) : null, context.organizationId ? resolveViewerSourceAccounts({ organizationId: context.organizationId, - userId: principal.userId, + userId: userId, connectors: scanned, }) : new Map(), diff --git a/apps/sim/lib/knowledge/application/sim-search.ts b/apps/sim/lib/knowledge/application/sim-search.ts index a8b1c94c5d8..e6940e16e02 100644 --- a/apps/sim/lib/knowledge/application/sim-search.ts +++ b/apps/sim/lib/knowledge/application/sim-search.ts @@ -21,6 +21,7 @@ import { } from '@/lib/core/resource-scope' import { resourceScopeCondition } from '@/lib/core/resource-scope.server' import { generateRequestId } from '@/lib/core/utils/request' +import type { CredentialGroupConnectionIntent } from '@/lib/credential-groups/oauth-intent' import { ensureWorkspaceAccountsGroup } from '@/lib/credential-groups/service' import { requireKnowledgeMemberAccessAvailable, @@ -65,6 +66,7 @@ export interface ConnectSimSearchConnectorInput extends ResourceOwner { /** Source settings identify a compatible configuration when creating or reusing a source. */ sourceConfig?: Record /** Correlates a direct provider authorization with the initiating Integrations tab. */ + connectionIntent?: CredentialGroupConnectionIntent oauthCompletionId?: string } @@ -370,6 +372,7 @@ export const connectSimSearchConnector = defineAuthorizedKnowledgeUseCase({ assertedWorkspaceId: workspaceId, assertedOrganizationId: context.organizationId, oauthCompletionId: input.oauthCompletionId, + connectionIntent: input.connectionIntent, }, request, }) diff --git a/apps/sim/lib/knowledge/application/slack-search/assistant.test.ts b/apps/sim/lib/knowledge/application/slack-search/assistant.test.ts index 14bb18b580b..d2f335c3e08 100644 --- a/apps/sim/lib/knowledge/application/slack-search/assistant.test.ts +++ b/apps/sim/lib/knowledge/application/slack-search/assistant.test.ts @@ -99,6 +99,7 @@ vi.mock('@/lib/copilot/request/session/abort', () => ({ startAbortPoller: () => 0, unregisterActiveStream: vi.fn(), })) +vi.mock('@/lib/slack-search/connections', () => ({ deliverSlackSearchConnections: vi.fn() })) vi.mock('@/lib/slack-search/assistant-stream', () => ({ SlackSearchAssistantStream: class { constructor(options: unknown) { @@ -188,40 +189,29 @@ describe('organization Assistant from Slack', () => { expect(m.title).not.toHaveBeenCalled() expect(m.run).not.toHaveBeenCalled() }) - it('persists source onboarding in private history without spending an Assistant run', async () => { + it('runs the Assistant with no indexed documents so it can list and connect integrations', async () => { m.sources.mockResolvedValueOnce({ hasSearchableDocuments: false }) await run() - expect(m.onboarding).toHaveBeenCalledWith( - principal, - expect.objectContaining({ reason: 'sources' }) - ) - expect(m.run).not.toHaveBeenCalled() - expect(m.createRun).not.toHaveBeenCalled() - expect(m.title).toHaveBeenCalledWith( - expect.objectContaining({ subjectUserId: 'member1', organizationId: 'org1' }), - expect.objectContaining({ job, signal: expect.any(AbortSignal) }) - ) + expect(m.onboarding).not.toHaveBeenCalled() + expect(m.run).toHaveBeenCalledOnce() + expect(m.createRun).toHaveBeenCalledOnce() expect(m.finalize).toHaveBeenCalledWith( - expect.objectContaining({ - assistantMessage: expect.objectContaining({ - content: expect.stringContaining('[Connect sources]'), - }), - }) + expect.objectContaining({ assistantMessage: expect.objectContaining({ content: 'Answer' }) }) ) }) - it('delivers onboarding while naming runs and awaits naming before releasing the worker', async () => { + it('answers while naming runs and awaits naming before releasing the worker', async () => { const naming = Promise.withResolvers() m.title.mockReturnValueOnce(naming.promise) m.sources.mockResolvedValueOnce({ hasSearchableDocuments: false }) const pending = run() - await vi.waitFor(() => expect(m.onboarding).toHaveBeenCalled()) + await vi.waitFor(() => expect(m.run).toHaveBeenCalled()) expect(m.release).not.toHaveBeenCalled() naming.resolve() await pending expect(m.release).toHaveBeenCalled() }) it.each(['missing response', 'database error'] as const)( - 'fails source onboarding when history persistence reports %s', + 'fails an answer when history persistence reports %s', async (outcome) => { m.sources.mockResolvedValueOnce({ hasSearchableDocuments: false }) if (outcome === 'missing response') { @@ -234,8 +224,8 @@ describe('organization Assistant from Slack', () => { ? 'Could not persist the Slack Assistant response' : 'History unavailable' ) - expect(m.onboarding).toHaveBeenCalledOnce() - expect(m.run).not.toHaveBeenCalled() + expect(m.onboarding).not.toHaveBeenCalled() + expect(m.run).toHaveBeenCalledOnce() expect(m.release).toHaveBeenCalledOnce() } ) diff --git a/apps/sim/lib/knowledge/application/slack-search/assistant.ts b/apps/sim/lib/knowledge/application/slack-search/assistant.ts index 8628b03be93..552a68ac710 100644 --- a/apps/sim/lib/knowledge/application/slack-search/assistant.ts +++ b/apps/sim/lib/knowledge/application/slack-search/assistant.ts @@ -43,13 +43,13 @@ import { } from '@/lib/knowledge/application/slack-search/identity' import { sendSlackSearchOnboarding } from '@/lib/knowledge/application/slack-search/onboarding' import { recordSlackSearchOutcome } from '@/lib/knowledge/application/slack-search/repository' -import { getSlackSearchSourceStatus } from '@/lib/knowledge/application/slack-search/source-status' import { generateSlackSearchChatTitle } from '@/lib/knowledge/application/slack-search/title' import { requireSlackSearchTurnLease, wasSlackSearchTurnStopped, } from '@/lib/knowledge/application/slack-search/turns' import { SlackSearchAssistantStream } from '@/lib/slack-search/assistant-stream' +import { deliverSlackSearchConnections } from '@/lib/slack-search/connections' import { SLACK_SEARCH_FAILED_ANSWER, SLACK_SEARCH_MAX_DURATION_SECONDS, @@ -190,101 +190,90 @@ export async function runSlackSearchAssistant( includeSecrets: false, }) registry = environmentContext.resolvedSecretTraceRegistry - const sources = await getSlackSearchSourceStatus.execute({ - principal: memberPrincipal(), - input: { organizationId: installation.organizationId }, + const executionId = generateId() + const run = await createRunSegment({ + executionId, + chatId: chat.id, + userId, + streamId: messageId, + model: chat.model, + status: 'active', }) - if (!sources.hasSearchableDocuments) { - await checkAccess() - const notice = await sendSlackSearchOnboarding(principal, { - job, - turnId, - leaseId, - email: sender.email, - reason: 'sources', - signal: controller.signal, - }) - const content = `${notice.text}\n\n[Connect sources](${notice.url})` - result = { - success: true, - content, - contentBlocks: [{ type: 'text', content, timestamp: Date.now() }], - toolCalls: [], - } - failed = false - } else { - const executionId = generateId() - const run = await createRunSegment({ - executionId, - chatId: chat.id, - userId, - streamId: messageId, - model: chat.model, - status: 'active', - }) - if (!run) throw new Error('Could not persist Assistant execution') - runId = run.id - const responseStream = new SlackSearchAssistantStream({ - token: secret.botToken, - channel: job.message.channelId, - threadTs: job.message.threadTs ?? job.message.messageTs, - slackUserId: job.message.userId, - controller, - registry: environmentContext.resolvedSecretTraceRegistry, - beforeDelivery: checkAccess, - beforeCleanup: checkAccess, - }) - stream = responseStream - const payload = await buildCopilotRequestPayload( - { - message: job.message.query, - userId, - userMessageId: messageId, + if (!run) throw new Error('Could not persist Assistant execution') + runId = run.id + const responseStream = new SlackSearchAssistantStream({ + token: secret.botToken, + channel: job.message.channelId, + threadTs: job.message.threadTs ?? job.message.messageTs, + slackUserId: job.message.userId, + controller, + registry: environmentContext.resolvedSecretTraceRegistry, + beforeDelivery: checkAccess, + beforeCleanup: checkAccess, + deliverConnections: (targets) => + deliverSlackSearchConnections({ + targets, + token: secret.botToken, organizationId: installation.organizationId, + userId, chatId: chat.id, - mode: 'assistant', - model: '', - }, - { selectedModel: '' } - ) - const billingAttribution = await resolveOrganizationBillingAttribution({ - actorUserId: userId, - organizationId: installation.organizationId, - }) - await responseStream.start() - result = await runHeadlessCopilotLifecycle(payload, { + turnId, + channel: job.message.channelId, + slackUserId: job.message.userId, + signal: controller.signal, + beforeDelivery: checkAccess, + }), + }) + stream = responseStream + const payload = await buildCopilotRequestPayload( + { + message: job.message.query, userId, + userMessageId: messageId, organizationId: installation.organizationId, chatId: chat.id, - executionId, - runId, - goRoute: '/api/mothership', - billingAttribution, - environmentContext, - resolvedSecretTraceRegistry: environmentContext.resolvedSecretTraceRegistry, - abortSignal: controller.signal, - timeout: SLACK_SEARCH_MAX_DURATION_SECONDS * 1000, - autoExecuteTools: true, - onEvent: async (event) => { - try { - await responseStream.onEvent(event) - } catch (error) { - controller.abort(error) - throw error - } - }, - }) - responseStream.assertHealthy() - controller.signal.throwIfAborted() - if (!result.success) { - await responseStream.finishWithError() - throw new Error('Organization Assistant did not complete') - } - await checkAccess() - await responseStream.finish(result) - failed = false - await recordSlackSearchOutcome(installation, 'success') + mode: 'assistant', + model: '', + }, + { selectedModel: '' } + ) + const billingAttribution = await resolveOrganizationBillingAttribution({ + actorUserId: userId, + organizationId: installation.organizationId, + }) + await responseStream.start() + result = await runHeadlessCopilotLifecycle(payload, { + userId, + organizationId: installation.organizationId, + chatId: chat.id, + executionId, + runId, + goRoute: '/api/mothership', + billingAttribution, + environmentContext, + resolvedSecretTraceRegistry: environmentContext.resolvedSecretTraceRegistry, + abortSignal: controller.signal, + timeout: SLACK_SEARCH_MAX_DURATION_SECONDS * 1000, + autoExecuteTools: true, + onEvent: async (event) => { + try { + await responseStream.onEvent(event) + } catch (error) { + controller.abort(error) + throw error + } + }, + }) + responseStream.assertHealthy() + controller.signal.throwIfAborted() + if (!result.success) { + await responseStream.finishWithError() + throw new Error('Organization Assistant did not complete') } + await checkAccess() + await responseStream.finish(result) + failed = false + await recordSlackSearchOutcome(installation, 'success') } catch (error) { failure = toError(error) controller.abort(error) diff --git a/apps/sim/lib/knowledge/application/slack-search/onboarding.test.ts b/apps/sim/lib/knowledge/application/slack-search/onboarding.test.ts index 7f85d50b285..89ef0204401 100644 --- a/apps/sim/lib/knowledge/application/slack-search/onboarding.test.ts +++ b/apps/sim/lib/knowledge/application/slack-search/onboarding.test.ts @@ -206,7 +206,7 @@ describe('Slack onboarding authorization and retry', () => { queueContext() m.sources.mockResolvedValueOnce({ hasSearchableDocuments: false }) expect(await get()).toEqual({ - status: 'needs_sources', + status: 'ready', organizationId: 'org1', question: job.message.query, isAdmin: false, @@ -215,11 +215,11 @@ describe('Slack onboarding authorization and retry', () => { expect(m.persist).not.toHaveBeenCalled() expect(m.dispatch).not.toHaveBeenCalled() }) - it('waits for accessible indexing before accepting a retry', async () => { + it('allows an authorized member to ask about integrations before indexing', async () => { queueContext() m.sources.mockResolvedValueOnce({ hasSearchableDocuments: false }) - await expect(retry()).rejects.toThrow('wait for indexing') - expect(m.persist).not.toHaveBeenCalled() + await expect(retry()).resolves.toEqual({ slackUrl: state.slackUrl }) + expect(m.persist).toHaveBeenCalledOnce() }) it('rechecks current capability permissions instead of trusting the link', async () => { queueContext() diff --git a/apps/sim/lib/knowledge/application/slack-search/onboarding.ts b/apps/sim/lib/knowledge/application/slack-search/onboarding.ts index 4cd4c7d79b1..5fd5eb22994 100644 --- a/apps/sim/lib/knowledge/application/slack-search/onboarding.ts +++ b/apps/sim/lib/knowledge/application/slack-search/onboarding.ts @@ -21,7 +21,6 @@ import { } from '@/lib/knowledge/application/slack-search/identity' import { dispatchSlackSearchTurn } from '@/lib/knowledge/application/slack-search/outbox' import { recordSlackSearchOutcome } from '@/lib/knowledge/application/slack-search/repository' -import { getSlackSearchSourceStatus } from '@/lib/knowledge/application/slack-search/source-status' import { persistSlackSearchTurn, requireSlackSearchTurnLease, @@ -297,12 +296,8 @@ async function resolveOnboarding(principal: Principal, token: string) { ) ) .limit(1) - const sources = await getSlackSearchSourceStatus.execute({ - principal, - input: { organizationId: context.installation.organizationId }, - }) const view: OnboardingReady = { - status: retried ? 'retried' : sources.hasSearchableDocuments ? 'ready' : 'needs_sources', + status: retried ? 'retried' : 'ready', organizationId: context.installation.organizationId, isAdmin: isOrgAdminRole(membership.role), question: job.message.query, @@ -334,11 +329,6 @@ export const retrySlackSearchOnboarding: OperationUseCase< const resolved = await resolveOnboarding(principal, input.token) if (!resolved.job || !('slackUrl' in resolved.view)) throw new OrchestrationError('forbidden', 'Complete your Sim account setup before retrying') - if (resolved.view.status === 'needs_sources') - throw new OrchestrationError( - 'validation', - 'Connect a source and wait for indexing before retrying' - ) let turnId = resolved.retryTurnId if (!turnId) { const now = Date.now() diff --git a/apps/sim/lib/knowledge/connectors/viewer-source-accounts.test.ts b/apps/sim/lib/knowledge/connectors/viewer-source-accounts.test.ts index 195511f6601..12182c73713 100644 --- a/apps/sim/lib/knowledge/connectors/viewer-source-accounts.test.ts +++ b/apps/sim/lib/knowledge/connectors/viewer-source-accounts.test.ts @@ -28,6 +28,7 @@ const account = { groupId: 'group-1', optionId: 'gmail-option', providerId: 'gmail', + status: 'active', } describe('personal source account projection', () => { @@ -46,10 +47,13 @@ describe('personal source account projection', () => { expect(isNull).toHaveBeenCalledWith(credentialGroup.workspaceId) expect(isNull).toHaveBeenCalledWith(credential.revokedAt) expect(inArray).toHaveBeenCalledWith(credential.managedOauthStatus, ['active', 'needs_reauth']) - expect(result.get(source.id)).toEqual([{ credentialId: 'mine', displayName: 'My Gmail' }]) + expect(result.get(source.id)).toEqual([ + { credentialId: 'mine', displayName: 'My Gmail', status: 'active' }, + ]) expect(dbChainMockFns.select).toHaveBeenCalledWith({ credentialId: credential.id, displayName: credential.displayName, + status: credential.managedOauthStatus, groupId: credentialGroup.id, optionId: credential.credentialGroupOptionId, providerId: credential.providerId, @@ -84,7 +88,7 @@ describe('personal source account projection', () => { expect(eq).toHaveBeenCalledWith(credential.type, 'managed_oauth') expect(eq).toHaveBeenCalledWith(credential.providerId, 'slack') expect(result.get('slack-source')).toEqual([ - { credentialId: 'slack-personal', displayName: 'My Gmail' }, + { credentialId: 'slack-personal', displayName: 'My Gmail', status: 'active' }, ]) }) diff --git a/apps/sim/lib/knowledge/connectors/viewer-source-accounts.ts b/apps/sim/lib/knowledge/connectors/viewer-source-accounts.ts index 84a501a0c75..2c9041e4acb 100644 --- a/apps/sim/lib/knowledge/connectors/viewer-source-accounts.ts +++ b/apps/sim/lib/knowledge/connectors/viewer-source-accounts.ts @@ -8,6 +8,7 @@ import { getConnectorMeta } from '@/connectors/registry' interface ViewerSourceAccount { credentialId: string displayName: string + status: 'active' | 'needs_reauth' } interface SourceAccountBinding { @@ -56,6 +57,7 @@ export async function resolveViewerSourceAccounts(input: { .select({ credentialId: credential.id, displayName: credential.displayName, + status: credential.managedOauthStatus, groupId: credentialGroup.id, optionId: credential.credentialGroupOptionId, providerId: credential.providerId, @@ -90,7 +92,11 @@ export async function resolveViewerSourceAccounts(input: { if (own.length) result.set( source.id, - own.map(({ credentialId, displayName }) => ({ credentialId, displayName })) + own.map(({ credentialId, displayName, status }) => { + if (status !== 'active' && status !== 'needs_reauth') + throw new Error('Invalid personal account status') + return { credentialId, displayName, status } + }) ) } return result diff --git a/apps/sim/lib/knowledge/search/connection-attempt.ts b/apps/sim/lib/knowledge/search/connection-attempt.ts new file mode 100644 index 00000000000..3eec8ee04c5 --- /dev/null +++ b/apps/sim/lib/knowledge/search/connection-attempt.ts @@ -0,0 +1,34 @@ +import { z } from 'zod' + +export const SEARCH_CONNECTION_ATTEMPT_MAX_AGE_MS = 10 * 60_000 +export const SEARCH_CONNECTION_ATTEMPT_EVENT = 'sim:search-connection-attempt' +const attemptSchema = z.object({ + completionId: z.string().uuid(), + requestedAt: z.number(), + connectorId: z.string().optional(), + credentialId: z.string().optional(), + status: z.enum(['pending', 'connected', 'failed']), + error: z.string().nullable(), +}) +export type SearchConnectionAttempt = z.infer + +/** Local presentation state, namespaced by the person, organization, and individual card. */ +export function searchConnectionAttemptKey( + organizationId: string, + userId: string, + controlId: string +) { + return `sim.search-connection.${encodeURIComponent(organizationId)}.${encodeURIComponent(userId)}.${encodeURIComponent(controlId)}` +} + +export function readSearchConnectionAttempt(key: string): SearchConnectionAttempt | null { + if (typeof window === 'undefined') return null + const value = window.localStorage.getItem(key) + if (!value) return null + return attemptSchema.parse(JSON.parse(value)) +} + +export function writeSearchConnectionAttempt(key: string, value: SearchConnectionAttempt) { + window.localStorage.setItem(key, JSON.stringify(value)) + window.dispatchEvent(new Event(SEARCH_CONNECTION_ATTEMPT_EVENT)) +} diff --git a/apps/sim/lib/knowledge/search/connection-target.test.ts b/apps/sim/lib/knowledge/search/connection-target.test.ts new file mode 100644 index 00000000000..323e81dcb26 --- /dev/null +++ b/apps/sim/lib/knowledge/search/connection-target.test.ts @@ -0,0 +1,38 @@ +/** @vitest-environment node */ +import { describe, expect, it } from 'vitest' +import { + parseSearchConnectionTargets, + searchConnectionPath, +} from '@/lib/knowledge/search/connection-target' + +const target = { + type: 'link', + provider: 'google-email', + connectorType: 'gmail', + connectorId: 'source', +} +describe('shared Search connection tags', () => { + it('parses complete arrays and deduplicates exact targets', () => { + expect( + parseSearchConnectionTargets(`${JSON.stringify([target, target])}`) + ).toEqual([target]) + }) + it.each([ + { ...target, value: 'https://evil.test' }, + { ...target, organizationId: 'other' }, + { ...target, connectorId: undefined, credentialId: 'account' }, + ])('rejects model URLs, scope and incomplete reconnects', (forged) => { + expect( + parseSearchConnectionTargets(`${JSON.stringify(forged)}`) + ).toEqual([]) + }) + it('holds partial and malformed tags until a complete validated target exists', () => { + expect(parseSearchConnectionTargets(`${JSON.stringify(target)}`)).toEqual([]) + expect(parseSearchConnectionTargets('{oops}')).toEqual([]) + }) + it('links to the existing organization page using only the selected IDs', () => { + expect(searchConnectionPath('org', target)).toBe( + '/o/org/integrations?connectorType=gmail&connectorId=source' + ) + }) +}) diff --git a/apps/sim/lib/knowledge/search/connection-target.ts b/apps/sim/lib/knowledge/search/connection-target.ts new file mode 100644 index 00000000000..fd00da116c8 --- /dev/null +++ b/apps/sim/lib/knowledge/search/connection-target.ts @@ -0,0 +1,52 @@ +import { z } from 'zod' +import { organizationRoutes } from '@/lib/navigation/paths' + +/** A requested personal Search connection; authority always comes from the current session. */ +export const searchConnectionTargetSchema = z + .object({ + type: z.literal('link'), + provider: z.string().trim().min(1).max(100), + connectorType: z.string().trim().min(1).max(100), + connectorId: z.string().min(1).max(200).optional(), + credentialId: z.string().min(1).max(128).optional(), + }) + .strict() + .refine((target) => !target.credentialId || Boolean(target.connectorId), { + message: 'A reconnect requires a configured source', + }) + +export type SearchConnectionTarget = z.infer + +/** Parses a single bounded Search card body; browser and Slack accept the same targets. */ +export function parseSearchConnectionBody(body: string): SearchConnectionTarget[] | null { + let value: unknown + try { + value = JSON.parse(body) + } catch { + return null + } + const parsed = z + .array(searchConnectionTargetSchema) + .min(1) + .max(10) + .safeParse(Array.isArray(value) ? value : [value]) + return parsed.success ? parsed.data : null +} + +/** Parses complete connection cards while withholding partial or malformed model output. */ +export function parseSearchConnectionTargets(text: string): SearchConnectionTarget[] { + const targets: SearchConnectionTarget[] = [] + for (const match of text.matchAll(/([\s\S]*?)<\/credential>/g)) { + for (const target of parseSearchConnectionBody(match[1]) ?? []) { + if (!targets.some((existing) => JSON.stringify(existing) === JSON.stringify(target))) + targets.push(target) + } + if (targets.length > 10) throw new Error('Too many requested Search connections') + } + return targets +} + +/** Carries an untrusted selection to the existing authenticated Integrations page, without OAuth state. */ +export function searchConnectionPath(organizationId: string, target: SearchConnectionTarget) { + return `${organizationRoutes(organizationId).integrations}?${new URLSearchParams({ connectorType: target.connectorType, ...(target.connectorId ? { connectorId: target.connectorId } : {}), ...(target.credentialId ? { credentialId: target.credentialId } : {}) })}` +} diff --git a/apps/sim/lib/slack-search/assistant-stream.test.ts b/apps/sim/lib/slack-search/assistant-stream.test.ts index e96c3cdbf6f..4fa5a9b4753 100644 --- a/apps/sim/lib/slack-search/assistant-stream.test.ts +++ b/apps/sim/lib/slack-search/assistant-stream.test.ts @@ -21,6 +21,10 @@ vi.mock('@/executor/utils/resolved-secret-content-projection', () => ({ projectResolvedSecretDiagnosticContent: api.project, })) +import type { + ToolCallStreamEvent, + ToolResultStreamEvent, +} from '@/lib/copilot/request/session/contract' import type { OrchestratorResult } from '@/lib/copilot/request/types' import { publicSlackAnswer, SlackSearchAssistantStream } from '@/lib/slack-search/assistant-stream' import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' @@ -55,7 +59,7 @@ function retrieval( }, } } -function setup() { +function setup(deliverConnections = vi.fn().mockResolvedValue(undefined)) { const controller = new AbortController() const beforeDelivery = vi.fn().mockResolvedValue(undefined) const beforeCleanup = vi.fn().mockResolvedValue(undefined) @@ -76,10 +80,192 @@ function setup() { registry, beforeDelivery, beforeCleanup, + deliverConnections, }), } } + +function toolCall(toolName = 'list_integrations', toolCallId = 'tool-1'): ToolCallStreamEvent { + return { + type: 'tool', + payload: { + phase: 'call', + toolName, + toolCallId, + executor: 'sim', + mode: 'sync', + status: 'executing', + }, + } +} + +function toolResult( + toolName = 'list_integrations', + toolCallId = 'tool-1', + success = true +): ToolResultStreamEvent { + return { + type: 'tool', + payload: { phase: 'result', toolName, toolCallId, executor: 'sim', mode: 'sync', success }, + } +} + +describe('Slack tool progress', () => { + it.each([ + ['list_integrations', 'Listing connected integrations…'], + ['search_workspace', 'Searching documents…'], + ['read_document', 'Reading documents…'], + ])('shows %s as a task and completes that same task once', async (name, title) => { + const { stream } = setup() + await stream.start() + await stream.onEvent(toolCall(name)) + await stream.onEvent(toolCall(name)) + await stream.onEvent(toolResult(name)) + await stream.onEvent(toolResult(name)) + await stream.finish(result) + const chunks = api.append.mock.calls.flatMap((call) => call[3]) + expect(chunks).toEqual([ + { type: 'task_update', id: expect.any(String), title, status: 'in_progress' }, + { type: 'task_update', id: chunks[0].id, title, status: 'complete' }, + ]) + expect(api.stop.mock.calls[0][6]).toEqual([]) + }) + + it('keeps parallel calls separate when their results arrive out of order', async () => { + const { stream } = setup() + await stream.start() + await stream.onEvent(toolCall('search_workspace', 'search-1')) + await stream.onEvent(toolCall('search_workspace', 'search-2')) + await stream.onEvent(toolResult('search_workspace', 'search-2')) + await stream.onEvent(toolResult('search_workspace', 'search-1')) + const chunks = api.append.mock.calls.flatMap((call) => call[3]) + expect(chunks[0].id).not.toBe(chunks[1].id) + expect(chunks[2]).toEqual({ ...chunks[1], status: 'complete' }) + expect(chunks[3]).toEqual({ ...chunks[0], status: 'complete' }) + }) + + it('withholds partial, hidden, internal, subagent, and unsupported tools', async () => { + const { stream } = setup() + await stream.start() + for (const attributes of [ + { partial: true }, + { status: 'generating' as const }, + { ui: { hidden: true } }, + { ui: { internal: true } }, + ]) { + const event = toolCall() + await stream.onEvent({ ...event, payload: { ...event.payload, ...attributes } }) + } + await stream.onEvent({ ...toolCall(), scope: { lane: 'subagent', agentId: 'private-agent' } }) + await stream.onEvent(toolCall('internal_tool')) + await stream.onEvent(toolResult()) + expect(api.append).not.toHaveBeenCalled() + await stream.onEvent(toolCall()) + expect(api.append).toHaveBeenCalledOnce() + }) + + it('reports failed tools without exposing arguments, account labels, or backend errors', async () => { + const { stream } = setup() + await stream.start() + const call = toolCall() + await stream.onEvent({ + ...call, + payload: { ...call.payload, arguments: { query: 'private argument' } }, + }) + const failed = toolResult('list_integrations', 'tool-1', false) + await stream.onEvent({ + ...failed, + payload: { + ...failed.payload, + error: 'private error', + output: { accountLabel: 'private account' }, + }, + }) + const chunks = api.append.mock.calls.flatMap((call) => call[3]) + expect(chunks[1]).toEqual({ ...chunks[0], status: 'error' }) + expect(JSON.stringify(chunks)).not.toContain('private') + }) + + it('marks unfinished tasks failed when the Assistant fails', async () => { + const { stream } = setup() + await stream.start() + await stream.onEvent(toolCall()) + await stream.finishWithError() + expect(api.stop.mock.calls[0][6]).toEqual([ + { ...api.append.mock.calls[0][3][0], status: 'error' }, + ]) + }) + + it('aborts an ambiguous progress send and cleans up once without replaying it', async () => { + const { stream, controller } = setup() + await stream.start() + api.append.mockRejectedValueOnce(new Error('progress response lost')) + await expect(stream.onEvent(toolCall())).rejects.toThrow('progress response lost') + expect(controller.signal.aborted).toBe(true) + await expect(stream.onEvent(toolCall())).rejects.toThrow('progress response lost') + await stream.terminateAfterFailure() + await stream.terminateAfterFailure() + expect(api.append).toHaveBeenCalledOnce() + expect(api.stop).toHaveBeenCalledOnce() + expect(api.stop.mock.calls[0][6]).toEqual([ + { ...api.append.mock.calls[0][3][0], status: 'error' }, + ]) + }) + + it('does not send task updates after cancellation or revoked delivery authority', async () => { + const { stream, controller, beforeDelivery } = setup() + await stream.start() + beforeDelivery.mockRejectedValueOnce(new Error('authority revoked')) + await expect(stream.onEvent(toolCall())).rejects.toThrow('authority revoked') + expect(controller.signal.aborted).toBe(true) + expect(api.append).not.toHaveBeenCalled() + const cancelled = setup() + await cancelled.stream.start() + cancelled.controller.abort(new Error('stopped')) + await expect(cancelled.stream.onEvent(toolCall())).rejects.toThrow('stopped') + expect(api.append).not.toHaveBeenCalled() + }) +}) + describe('Slack Assistant delivery', () => { + it('withholds split connection tags, delivers validated controls, and leaves a visible next step', async () => { + const deliver = vi.fn().mockResolvedValue(undefined) + const { stream } = setup(deliver) + await stream.start() + const target = { type: 'link', provider: 'google-email', connectorType: 'gmail' } + for (const text of [ + 'Connect Gmail. ${JSON.stringify(target).slice(0, 10)}`, + `${JSON.stringify(target).slice(10)}`, + ]) { + await stream.onEvent({ type: 'text', payload: { channel: 'assistant', text } }) + } + await stream.finish(result) + expect(deliver).toHaveBeenCalledExactlyOnceWith([target]) + expect(deliveredText()).toContain('connection buttons in our DM') + expect(deliveredText()).not.toMatch(/credential|connectorType|google-email/) + }) + it('aborts on connection-button delivery failure without a successful finish', async () => { + const deliver = vi.fn().mockRejectedValue(new Error('ephemeral delivery failed')) + const { stream, controller } = setup(deliver) + await stream.start() + await stream.onEvent({ + type: 'text', + payload: { + channel: 'assistant', + text: '{"type":"link","provider":"gmail","connectorType":"gmail"}', + }, + }) + await expect(stream.finish(result)).rejects.toThrow('ephemeral delivery failed') + expect(controller.signal.aborted).toBe(true) + expect(api.stop).not.toHaveBeenCalled() + }) + it('never exposes a partial terminal tag or model-authored connection URL', () => { + expect(publicSlackAnswer('Next {oops}', true)).toBe( + 'Connect here ' + ) + }) it('streams only main public answer text and preserves the original thread', async () => { const { stream } = setup() await stream.start() @@ -210,7 +396,8 @@ describe('Slack Assistant delivery', () => { type: 'section', text: { type: 'plain_text', text: 'I couldn’t complete this search. Please try again.' }, }, - ] + ], + [] ) expect(api.append).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/slack-search/assistant-stream.ts b/apps/sim/lib/slack-search/assistant-stream.ts index 79485fea134..201e1270f99 100644 --- a/apps/sim/lib/slack-search/assistant-stream.ts +++ b/apps/sim/lib/slack-search/assistant-stream.ts @@ -1,4 +1,5 @@ import { toError } from '@sim/utils/errors' +import { generateId } from '@sim/utils/id' import { truncate } from '@sim/utils/string' import { collectRetrievalCitationEvidence, @@ -6,11 +7,20 @@ import { type RetrievalCitationBlock, } from '@/lib/copilot/chat/citation-evidence' import { redactSensitiveContent } from '@/lib/copilot/chat/sim-key-redaction' -import type { StreamEvent } from '@/lib/copilot/request/session/contract' +import type { + StreamEvent, + ToolCallStreamEvent, + ToolResultStreamEvent, +} from '@/lib/copilot/request/session/contract' import type { OrchestratorResult } from '@/lib/copilot/request/types' +import { + parseSearchConnectionTargets, + type SearchConnectionTarget, +} from '@/lib/knowledge/search/connection-target' import { SLACK_SEARCH_FAILED_ANSWER } from '@/lib/slack-search/constants' import { appendSlackAgentStream, + type SlackStreamChunk, setSlackAgentSessionStatus, startSlackAgentStream, stopSlackAgentStream, @@ -54,7 +64,7 @@ export function publicSlackAnswer( function publicSlackText(value: string): string { return value .replace(/\[([^\]]*)\]\([^)]*\)/g, '$1') - .replace(/<[^>]*>/g, '') + .replace(/<[^>]*(?:>|$)/g, '') .replace(/(?:https?:\/\/|www\.)[^\s<>]+/gi, '') .replaceAll('&', '&') .replaceAll('<', '<') @@ -89,12 +99,21 @@ interface AssistantStreamOptions { registry: ResolvedSecretTraceRegistry beforeDelivery: () => Promise beforeCleanup: (signal: AbortSignal) => Promise + deliverConnections?: (targets: SearchConnectionTarget[]) => Promise } const FAILURE_BLOCKS: Record[] = [ { type: 'section', text: { type: 'plain_text', text: SLACK_SEARCH_FAILED_ANSWER } }, ] +const TOOL_PROGRESS_TITLES = new Map([ + ['list_integrations', 'Listing connected integrations…'], + ['search_workspace', 'Searching documents…'], + ['read_document', 'Reading documents…'], +]) + +type ToolProgress = Extract + /** Serial delivery through the same provider primitives as Slack blocks; ambiguous sends are terminal. */ export class SlackSearchAssistantStream { private stream?: { channel: string; ts: string } @@ -106,6 +125,7 @@ export class SlackSearchAssistantStream { private closeAttempted = false private separateNextText = false private evidence = new Map>() + private toolProgress = new Map() constructor(private readonly options: AssistantStreamOptions) {} private async deliver(action: () => Promise) { @@ -154,7 +174,15 @@ export class SlackSearchAssistantStream { }, ]) } - if (event.type === 'tool' && !event.scope) this.separateNextText = true + if (event.type === 'tool' && !event.scope) { + this.separateNextText = true + if ( + 'phase' in event.payload && + (event.payload.phase === 'call' || event.payload.phase === 'result') + ) { + await this.updateToolProgress(event.payload) + } + } if (event.type !== 'text' || event.payload.channel !== 'assistant' || event.scope) return if (this.separateNextText && this.text) this.text += '\n\n' this.separateNextText = false @@ -163,6 +191,57 @@ export class SlackSearchAssistantStream { if (Date.now() - this.lastSentAt >= 750) await this.flush(false) } + /** Only static labels reach Slack; arguments, account details, and backend errors stay private. */ + private async updateToolProgress( + payload: ToolCallStreamEvent['payload'] | ToolResultStreamEvent['payload'] + ) { + const title = TOOL_PROGRESS_TITLES.get(payload.toolName) + if (!title) return + const existing = this.toolProgress.get(payload.toolCallId) + let chunk: ToolProgress + if (payload.phase === 'call') { + if ( + existing || + payload.partial || + payload.ui?.hidden || + payload.ui?.internal || + (payload.status !== undefined && payload.status !== 'executing') + ) + return + chunk = { type: 'task_update', id: generateId(), title, status: 'in_progress' } + this.toolProgress.set(payload.toolCallId, { toolName: payload.toolName, chunk }) + } else { + if (!existing || existing.chunk.status !== 'in_progress') return + if (existing.toolName !== payload.toolName) + throw new Error('Slack tool progress identity changed') + chunk = { + ...existing.chunk, + status: + payload.success && (!payload.status || payload.status === 'success') + ? 'complete' + : 'error', + } + } + await this.deliver(async () => { + if (!this.stream || this.closed) throw new Error('Slack stream is not active') + await appendSlackAgentStream( + this.options.token, + this.stream.channel, + this.stream.ts, + [chunk], + this.options.controller.signal + ) + }) + this.toolProgress.set(payload.toolCallId, { toolName: payload.toolName, chunk }) + } + + /** Finalize interrupted tasks in the single stop request, including ambiguous progress sends. */ + private interruptedToolProgress(): ToolProgress[] { + return [...this.toolProgress.values()] + .filter(({ chunk }) => chunk.status === 'in_progress') + .map(({ chunk }) => ({ ...chunk, status: 'error' })) + } + private collectSources(blocks: readonly RetrievalCitationBlock[]) { for (const [id, source] of collectRetrievalCitationEvidence(blocks)) { if (!this.evidence.has(id)) this.evidence.set(id, source) @@ -221,6 +300,21 @@ export class SlackSearchAssistantStream { async finish(result: OrchestratorResult) { if (this.failure) throw this.failure this.collectSources(result.contentBlocks) + const projection = projectResolvedSecretDiagnosticContent( + this.text, + this.options.registry, + 512_000 + ) + if (!projection.safe || typeof projection.value !== 'string') + throw new Error('Connection controls could not be safely projected') + const targets = parseSearchConnectionTargets(redactSensitiveContent(projection.value)) + if (targets.length) { + if (!this.options.deliverConnections) + throw new Error('Search connection delivery is unavailable') + await this.deliver(() => this.options.deliverConnections!(targets)) + this.text += + '\n\nUse the connection buttons in our DM, then reply here when you’re ready to continue.' + } await this.flush(true) await this.close([]) } @@ -243,7 +337,8 @@ export class SlackSearchAssistantStream { this.stream.ts, 'active', signal, - FAILURE_BLOCKS + FAILURE_BLOCKS, + this.interruptedToolProgress() ) this.closed = true } @@ -258,7 +353,8 @@ export class SlackSearchAssistantStream { this.stream.ts, 'active', this.options.controller.signal, - blocks + blocks, + this.interruptedToolProgress() ) this.closed = true }) diff --git a/apps/sim/lib/slack-search/connections.test.ts b/apps/sim/lib/slack-search/connections.test.ts new file mode 100644 index 00000000000..2a95cdec8bb --- /dev/null +++ b/apps/sim/lib/slack-search/connections.test.ts @@ -0,0 +1,79 @@ +/** @vitest-environment node */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const m = vi.hoisted(() => ({ read: vi.fn(), send: vi.fn(), before: vi.fn(), origin: vi.fn() })) +vi.mock('@/lib/copilot/application/execute-knowledge-use-case', () => ({ + executeCopilotOrganizationKnowledgeUseCase: m.read, +})) +vi.mock('@/lib/knowledge/application/personal-search-integrations', () => ({ + resolvePersonalSearchConnection: {}, +})) +vi.mock('@/lib/core/utils/urls', () => ({ getBaseUrl: m.origin })) +vi.mock('@/lib/internal/slack/client', () => ({ requestSlackApi: m.send })) + +import { deliverSlackSearchConnections } from '@/lib/slack-search/connections' + +const target = { + type: 'link', + provider: 'google-email', + connectorType: 'gmail', + connectorId: 'source', +} as const +const input = { + targets: [target], + token: 'token', + organizationId: 'org', + userId: 'person', + chatId: 'chat', + turnId: 'turn', + channel: 'D1', + slackUserId: 'U1', + signal: new AbortController().signal, + beforeDelivery: m.before, +} +beforeEach(() => { + vi.clearAllMocks() + m.read.mockResolvedValue({ name: 'Gmail', target }) + m.send.mockResolvedValue({ status: 200, data: { ok: true } }) + m.origin.mockReturnValue('https://preview.example.test') +}) +describe('Slack connection controls', () => { + it.each(['https://staging.example.test', 'https://preview.example.test'])( + 'uses configured origin %s and private ephemeral delivery', + async (origin) => { + m.origin.mockReturnValue(origin) + await deliverSlackSearchConnections(input) + expect(m.read).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'person', organizationId: 'org', chatId: 'chat' }), + {}, + { organizationId: 'org', target } + ) + expect(m.send).toHaveBeenCalledWith( + expect.objectContaining({ + method: 'chat.postEphemeral', + body: expect.objectContaining({ channel: 'D1', user: 'U1' }), + }) + ) + const button = m.send.mock.calls[0][0].body.blocks[1].elements[0] + expect(button.url).toBe(`${origin}/o/org/integrations?connectorType=gmail&connectorId=source`) + expect(button.text.text).toBe('Connect Gmail') + expect(m.before).toHaveBeenCalledTimes(2) + } + ) + it('never publishes personal account details to a channel', async () => { + await expect(deliverSlackSearchConnections({ ...input, channel: 'C1' })).rejects.toThrow( + 'require a Slack DM' + ) + expect(m.send).not.toHaveBeenCalled() + }) + it('fails before sending when current target authority has changed', async () => { + m.read.mockRejectedValue(new Error('target revoked')) + await expect(deliverSlackSearchConnections(input)).rejects.toThrow('target revoked') + expect(m.send).not.toHaveBeenCalled() + }) + it('fails on an ambiguous send without replaying or changing delivery mechanisms', async () => { + m.send.mockRejectedValue(new Error('connection reset')) + await expect(deliverSlackSearchConnections(input)).rejects.toThrow('connection reset') + expect(m.send).toHaveBeenCalledOnce() + }) +}) diff --git a/apps/sim/lib/slack-search/connections.ts b/apps/sim/lib/slack-search/connections.ts new file mode 100644 index 00000000000..b54df5679b2 --- /dev/null +++ b/apps/sim/lib/slack-search/connections.ts @@ -0,0 +1,77 @@ +import { executeCopilotOrganizationKnowledgeUseCase } from '@/lib/copilot/application/execute-knowledge-use-case' +import { getBaseUrl } from '@/lib/core/utils/urls' +import { requestSlackApi } from '@/lib/internal/slack/client' +import { resolvePersonalSearchConnection } from '@/lib/knowledge/application/personal-search-integrations' +import { + type SearchConnectionTarget, + searchConnectionPath, +} from '@/lib/knowledge/search/connection-target' + +interface SlackSearchConnectionsInput { + targets: SearchConnectionTarget[] + token: string + organizationId: string + userId: string + chatId: string + turnId: string + channel: string + slackUserId: string + signal: AbortSignal + beforeDelivery: () => Promise +} + +/** Validates each requested personal target before sending server-authored, ephemeral buttons. */ +export async function deliverSlackSearchConnections(input: SlackSearchConnectionsInput) { + if (!input.channel.startsWith('D')) throw new Error('Personal connections require a Slack DM') + const elements: Record[] = [] + for (const [index, target] of input.targets.entries()) { + await input.beforeDelivery() + const selected = await executeCopilotOrganizationKnowledgeUseCase( + { + userId: input.userId, + organizationId: input.organizationId, + chatId: input.chatId, + toolCallId: `${input.turnId}:connection:${index}`, + copilotToolExecution: true, + requestMode: 'assistant', + }, + resolvePersonalSearchConnection, + { organizationId: input.organizationId, target } + ) + elements.push({ + type: 'button', + action_id: `search_connect_${index}`, + text: { + type: 'plain_text', + text: `${target.credentialId ? 'Reconnect' : 'Connect'} ${selected.name}`, + }, + url: new URL(searchConnectionPath(input.organizationId, selected.target), getBaseUrl()).href, + }) + } + await input.beforeDelivery() + input.signal.throwIfAborted() + const blocks: Record[] = [ + { + type: 'section', + text: { + type: 'plain_text', + text: 'Connect your accounts in Sim, then reply in this thread to continue.', + }, + }, + ] + for (let offset = 0; offset < elements.length; offset += 5) + blocks.push({ type: 'actions', elements: elements.slice(offset, offset + 5) }) + const response = await requestSlackApi({ + accessToken: input.token, + method: 'chat.postEphemeral', + signal: input.signal, + body: { + channel: input.channel, + user: input.slackUserId, + text: 'Connect your sources in Sim', + blocks, + }, + }) + if (response.status !== 200 || response.data.ok !== true) + throw new Error('Could not deliver Search connection buttons') +} diff --git a/apps/sim/lib/webhooks/slack-agent-api.test.ts b/apps/sim/lib/webhooks/slack-agent-api.test.ts index 833816bc9e5..df0634696c9 100644 --- a/apps/sim/lib/webhooks/slack-agent-api.test.ts +++ b/apps/sim/lib/webhooks/slack-agent-api.test.ts @@ -85,6 +85,28 @@ describe('Slack agent API transport', () => { }) }) + it('finishes task updates in the same stop request as the final blocks', async () => { + const fetchMock = vi.fn().mockResolvedValue(new Response(JSON.stringify({ ok: true }))) + vi.stubGlobal('fetch', fetchMock) + const chunks = [ + { + type: 'task_update' as const, + id: 'task-1', + title: 'Searching documents…', + status: 'error' as const, + }, + ] + const blocks = [{ type: 'section', text: { type: 'plain_text', text: 'Please try again.' } }] + await stopSlackAgentStream('xoxb-test', 'D1', '101.2', 'active', undefined, blocks, chunks) + expect(JSON.parse(fetchMock.mock.calls[0][1].body)).toEqual({ + channel: 'D1', + ts: '101.2', + session_status: 'active', + blocks, + chunks, + }) + }) + it('sets the human initiator when creating a processing session', async () => { const fetchMock = vi .fn() diff --git a/apps/sim/lib/webhooks/slack-agent-api.ts b/apps/sim/lib/webhooks/slack-agent-api.ts index ed52ade4332..8e4f41392db 100644 --- a/apps/sim/lib/webhooks/slack-agent-api.ts +++ b/apps/sim/lib/webhooks/slack-agent-api.ts @@ -148,12 +148,19 @@ export async function stopSlackAgentStream( ts: string, sessionStatus: 'active' | 'processing' | 'suspended', signal?: AbortSignal, - blocks?: Record[] + blocks?: Record[], + chunks?: SlackStreamChunk[] ): Promise { await callSlackAgentApi( 'chat.stopStream', token, - { channel, ts, session_status: sessionStatus, ...(blocks?.length ? { blocks } : {}) }, + { + channel, + ts, + session_status: sessionStatus, + ...(blocks?.length ? { blocks } : {}), + ...(chunks?.length ? { chunks } : {}), + }, signal ) } diff --git a/apps/sim/tools/pitchbook/pitchbook.test.ts b/apps/sim/tools/pitchbook/pitchbook.test.ts index b816a40f217..7e280a13342 100644 --- a/apps/sim/tools/pitchbook/pitchbook.test.ts +++ b/apps/sim/tools/pitchbook/pitchbook.test.ts @@ -1,8 +1,11 @@ /** * @vitest-environment node */ +import { inputValidationMock, inputValidationMockFns } from '@sim/testing' import { describe, expect, it, vi } from 'vitest' +vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock) + /** * Only this service's configs are needed; the full registry is ~6,000 modules. * Registration is asserted through the generated `@/tools/tool-ids`. @@ -484,28 +487,42 @@ describe('pitchbook error extraction', () => { * not appear anywhere in the tool result, message or retained body. */ it('keeps the rejected key out of the whole failed tool result', async () => { - const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValue( - new Response( - JSON.stringify({ - reason: 'UNAUTHORIZED', - message: `Active API key ${SUBMITTED_KEY} not found`, - }), - { status: 401, headers: { 'content-type': 'application/json' } } - ) + const response = new Response( + JSON.stringify({ + reason: 'UNAUTHORIZED', + message: `Active API key ${SUBMITTED_KEY} not found`, + }), + { status: 401, headers: { 'content-type': 'application/json' } } ) + inputValidationMockFns.mockValidateUrlWithDNS.mockResolvedValueOnce({ + isValid: true, + resolvedIP: '93.184.216.34', + }) + inputValidationMockFns.mockSecureFetchWithPinnedIP.mockResolvedValueOnce({ + ok: response.ok, + status: response.status, + statusText: response.statusText, + headers: { + get: (name: string) => response.headers.get(name), + toRecord: () => Object.fromEntries(response.headers.entries()), + }, + body: response.body, + text: () => response.text(), + json: () => response.json(), + arrayBuffer: () => response.arrayBuffer(), + }) - try { - const result = await executeTool('pitchbook_company_bio', { - apiKey: SUBMITTED_KEY, - pbId: '10618-03', - }) - - expect(result.success).toBe(false) - expect(JSON.stringify(result.output ?? {})).not.toContain(SUBMITTED_KEY) - expect(result.error ?? '').not.toContain(SUBMITTED_KEY) - } finally { - fetchSpy.mockRestore() - } + const result = await executeTool('pitchbook_company_bio', { + apiKey: SUBMITTED_KEY, + pbId: '10618-03', + }) + + expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).toHaveBeenCalledOnce() + expect(result.success).toBe(false) + expect(result.error).toBe( + 'PitchBook rejected the API key. Check that the key is active and has API access.' + ) + expect(JSON.stringify(result)).not.toContain(SUBMITTED_KEY) }) it('routes every pitchbook tool through the scrubbing extractor', () => {