feat(search-mcp): persist client-attributed tool activity

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/organization-mcp-search.integration.ts

@@ -26,6 +26,7 @@ import {
   oauthConsent,
   organization,
   organizationSearchIntegration,
+  organizationSearchMcpInvocation,
   rateLimitBucket,
   user,
   workspace,
@@ -35,9 +36,19 @@ import { generateId } from '@sim/utils/id'
 import { isPlainRecord } from '@sim/utils/object'
 import { and, eq, inArray } from 'drizzle-orm'
 import { NextRequest } from 'next/server'
-import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
 
-const fixtures = vi.hoisted(() => ({ storageRoot: '' }))
+const fixtures = vi.hoisted(() => ({
+  storageRoot: '',
+  afterResponse: [] as Array<() => Promise<void>>,
+}))
+vi.mock('@/lib/core/utils/after-response', () => ({
+  afterResponse: (task: () => Promise<void>) => fixtures.afterResponse.push(task),
+}))
+
+async function flushAfterResponse() {
+  for (const task of fixtures.afterResponse.splice(0)) await task()
+}
 vi.mock('@/lib/uploads/core/setup.server', () => ({
   get UPLOAD_DIR_SERVER() {
     return fixtures.storageRoot
@@ -401,6 +412,8 @@ describe('organization Search MCP with real ingestion and current access', () =>
     bobOAuth = await connect(OAUTH_ACCESS_TOKEN_PREFIX + oauthTokens.bob, true)
   })
 
+  afterEach(flushAfterResponse)
+
   afterAll(async () => {
     await Promise.all(clients.map((client) => client.close()))
     await db.delete(oauthClient).where(eq(oauthClient.clientId, oauthClientId))
@@ -508,6 +521,72 @@ describe('organization Search MCP with real ingestion and current access', () =>
     expect(await applicationSearch(bobPrincipal)).toEqual([])
   })
 
+  it('persists content-free per-client tool outcomes separately from search counters', async () => {
+    await db
+      .delete(organizationSearchMcpInvocation)
+      .where(eq(organizationSearchMcpInvocation.organizationId, organizationId))
+    await db
+      .update(oauthClient)
+      .set({ name: 'MCP fixture client' })
+      .where(eq(oauthClient.clientId, oauthClientId))
+    await aliceOAuth.listTools()
+    expect(fixtures.afterResponse).toHaveLength(0)
+    await search(aliceOAuth)
+    await value(aliceOAuth, 'read_document', { documentId })
+    expect((await call(bob, 'read_document', { documentId })).isError).toBe(true)
+    expect(fixtures.afterResponse).toHaveLength(3)
+    await flushAfterResponse()
+    const rows = await db
+      .select()
+      .from(organizationSearchMcpInvocation)
+      .where(eq(organizationSearchMcpInvocation.organizationId, organizationId))
+      .orderBy(organizationSearchMcpInvocation.createdAt)
+      .limit(10)
+    expect(rows).toHaveLength(3)
+    expect(rows).toMatchObject([
+      {
+        organizationId,
+        userId: aliceId,
+        authKind: 'oauth_access_token',
+        oauthClientId,
+        clientName: 'MCP fixture client',
+        toolName: 'search',
+        outcome: 'success',
+      },
+      {
+        organizationId,
+        userId: aliceId,
+        authKind: 'oauth_access_token',
+        oauthClientId,
+        clientName: 'MCP fixture client',
+        toolName: 'read_document',
+        outcome: 'success',
+      },
+      {
+        organizationId,
+        userId: bobId,
+        authKind: 'personal_api_key',
+        oauthClientId: null,
+        clientName: null,
+        toolName: 'read_document',
+        outcome: 'error',
+      },
+    ])
+    expect(rows.every((row) => row.durationMs >= 0)).toBe(true)
+    expect(JSON.stringify(rows)).not.toContain(documentId)
+    expect(JSON.stringify(rows)).not.toContain(oauthTokens.alice)
+    await db
+      .update(oauthClient)
+      .set({ name: 'Renamed client' })
+      .where(eq(oauthClient.clientId, oauthClientId))
+    const [historical] = await db
+      .select()
+      .from(organizationSearchMcpInvocation)
+      .where(eq(organizationSearchMcpInvocation.id, rows[0].id))
+      .limit(1)
+    expect(historical.clientName).toBe('MCP fixture client')
+  })
+
   it('enforces current document and organization access on Search OAuth clients', async () => {
     expect((await aliceOAuth.listTools()).tools).toHaveLength(3)
     expect(await search(aliceOAuth)).toEqual(await search(alice))

apps/sim/lib/knowledge/mcp/activity.test.ts

@@ -0,0 +1,86 @@
+/** @vitest-environment node */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  values: vi.fn(),
+  insert: vi.fn(),
+  execute: vi.fn(),
+  transaction: vi.fn(),
+}))
+vi.mock('@sim/db', () => ({ db: { transaction: mocks.transaction } }))
+
+import {
+  recordOrganizationSearchMcpActivity,
+  type SearchMcpActivityInput,
+} from '@/lib/knowledge/mcp/activity'
+
+const activity: SearchMcpActivityInput = {
+  organizationId: 'org',
+  userId: 'actor',
+  authKind: 'personal_api_key',
+  oauthClientId: null,
+  toolName: 'read_document',
+  outcome: 'success',
+  durationMs: 42,
+  createdAt: new Date('2026-01-01T00:00:00Z'),
+}
+
+beforeEach(() => {
+  vi.clearAllMocks()
+  mocks.insert.mockReturnValue({ values: mocks.values })
+  mocks.values.mockResolvedValue(undefined)
+  mocks.execute.mockResolvedValue(undefined)
+  mocks.transaction.mockImplementation((callback) =>
+    callback({ execute: mocks.execute, insert: mocks.insert })
+  )
+})
+
+describe('persistent MCP activity', () => {
+  it('stores an API-key call without inventing an application name', async () => {
+    await recordOrganizationSearchMcpActivity(activity)
+    expect(mocks.values).toHaveBeenCalledExactlyOnceWith({
+      id: expect.any(String),
+      ...activity,
+      clientName: null,
+    })
+  })
+
+  it('only persists the allowlisted metadata when extra content is present', async () => {
+    const input = {
+      ...activity,
+      query: 'private question',
+      content: 'private document',
+      token: 'private token',
+    }
+    await recordOrganizationSearchMcpActivity(input)
+    expect(mocks.values).toHaveBeenCalledExactlyOnceWith({
+      id: expect.any(String),
+      ...activity,
+      clientName: null,
+    })
+  })
+
+  it('sets the transaction deadline before attempting the insert', async () => {
+    const ready = Promise.withResolvers<void>()
+    mocks.execute.mockReturnValueOnce(ready.promise)
+    const recording = recordOrganizationSearchMcpActivity(activity)
+    expect(mocks.insert).not.toHaveBeenCalled()
+    expect(JSON.stringify(mocks.execute.mock.calls[0])).toContain(
+      "SET LOCAL statement_timeout = '2s'"
+    )
+    ready.resolve()
+    await recording
+    expect(mocks.insert).toHaveBeenCalledOnce()
+  })
+
+  it('does not insert when the deadline could not be established', async () => {
+    mocks.execute.mockRejectedValueOnce(new Error('unavailable'))
+    await expect(recordOrganizationSearchMcpActivity(activity)).resolves.toBeUndefined()
+    expect(mocks.insert).not.toHaveBeenCalled()
+  })
+
+  it('does not propagate storage failures into the request lifecycle', async () => {
+    mocks.values.mockRejectedValueOnce(new Error('offline'))
+    await expect(recordOrganizationSearchMcpActivity(activity)).resolves.toBeUndefined()
+  })
+})

apps/sim/lib/knowledge/mcp/activity.ts

@@ -0,0 +1,49 @@
+import { db } from '@sim/db'
+import { oauthClient, organizationSearchMcpInvocation } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+import { generateId } from '@sim/utils/id'
+import { sql } from 'drizzle-orm'
+
+const logger = createLogger('OrganizationSearchMcpActivity')
+
+export type SearchMcpActivityInput = Pick<
+  typeof organizationSearchMcpInvocation.$inferInsert,
+  | 'organizationId'
+  | 'userId'
+  | 'authKind'
+  | 'oauthClientId'
+  | 'toolName'
+  | 'outcome'
+  | 'durationMs'
+  | 'createdAt'
+>
+
+/** Stores content-free metadata from an admitted MCP request, independently of tool success. */
+export async function recordOrganizationSearchMcpActivity(
+  input: SearchMcpActivityInput
+): Promise<void> {
+  try {
+    await db.transaction(async (tx) => {
+      await tx.execute(sql`SET LOCAL statement_timeout = '2s'`)
+      await tx.insert(organizationSearchMcpInvocation).values({
+        id: generateId(),
+        organizationId: input.organizationId,
+        userId: input.userId,
+        authKind: input.authKind,
+        oauthClientId: input.oauthClientId,
+        clientName: input.oauthClientId
+          ? sql`(SELECT left(${oauthClient.name}, 256) FROM ${oauthClient} WHERE ${oauthClient.clientId} = ${input.oauthClientId})`
+          : null,
+        toolName: input.toolName,
+        outcome: input.outcome,
+        durationMs: input.durationMs,
+        createdAt: input.createdAt,
+      })
+    })
+  } catch (error) {
+    logger.warn('Failed to record organization Search MCP activity', {
+      error: getErrorMessage(error),
+    })
+  }
+}

apps/sim/lib/knowledge/mcp/server.test.ts

@@ -1,5 +1,6 @@
 /** @vitest-environment node */
 import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import { createMockLogger } from '@sim/testing'
 import { NextRequest } from 'next/server'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -12,6 +13,16 @@ const mocks = vi.hoisted(() => ({
   read: vi.fn(),
   chat: vi.fn(),
   rateLimit: vi.fn(),
+  info: vi.fn(),
+  afterResponse: vi.fn<(task: () => Promise<void>) => void>(),
+  recordActivity: vi.fn(),
+}))
+vi.mock('@/lib/core/utils/after-response', () => ({ afterResponse: mocks.afterResponse }))
+vi.mock('@/lib/knowledge/mcp/activity', () => ({
+  recordOrganizationSearchMcpActivity: mocks.recordActivity,
+}))
+vi.mock('@sim/logger', () => ({
+  createLogger: () => ({ ...createMockLogger(), info: mocks.info }),
 }))
 vi.mock('@modelcontextprotocol/sdk/server/mcp.js', () => ({
   McpServer: class {
@@ -314,3 +325,115 @@ describe('organization chat', () => {
     })
   })
 })
+
+describe('MCP tool completion records', () => {
+  it('schedules only metadata after the response, independently of analytics storage latency', async () => {
+    createKnowledgeMcpServer({
+      organizationId: 'org-1',
+      searchIndexId: 'index-1',
+      request,
+      auth: {
+        ...auth,
+        keyType: 'oauth_access_token',
+        principal: {
+          kind: 'oauth_access_token',
+          userId: 'oauth-person',
+          clientId: 'registered-client',
+          tokenId: 'private-token-id',
+          scopes: ['search:read'],
+          expiresAt: new Date(Date.now() + 60000),
+        },
+      },
+    })
+    const response = await call('search', { query: 'private question' })
+    expect(response.isError).not.toBe(true)
+    expect(mocks.recordActivity).not.toHaveBeenCalled()
+    expect(mocks.afterResponse).toHaveBeenCalledOnce()
+    await mocks.afterResponse.mock.calls[0][0]()
+    expect(mocks.recordActivity).toHaveBeenCalledExactlyOnceWith({
+      organizationId: 'org-1',
+      userId: 'oauth-person',
+      authKind: 'oauth_access_token',
+      oauthClientId: 'registered-client',
+      toolName: 'search',
+      outcome: 'success',
+      durationMs: expect.any(Number),
+      createdAt: expect.any(Date),
+    })
+  })
+
+  it.each([
+    ['search', { query: 'private query', topK: 10 }, 'knowledge.search'],
+    ['read_document', { documentId: 'doc-1' }, 'knowledge.documents.read'],
+    ['chat', { query: 'private question' }, 'knowledge.chat'],
+  ] as const)('records one content-free completion for %s', async (toolName, input, operation) => {
+    create()
+    await call(toolName, input)
+    expect(mocks.info).toHaveBeenCalledExactlyOnceWith('Knowledge MCP tool completed', {
+      toolName,
+      operation,
+      organizationId: 'org-1',
+      userId: 'person-1',
+      outcome: 'success',
+      durationMs: expect.any(Number),
+    })
+  })
+
+  it('records a returned tool error as an error even though the HTTP transport can succeed', async () => {
+    create()
+    const result = await call('read_document', {})
+    expect(result.isError).toBe(true)
+    expect(mocks.info).toHaveBeenCalledExactlyOnceWith(
+      'Knowledge MCP tool completed',
+      expect.objectContaining({ toolName: 'read_document', outcome: 'error' })
+    )
+  })
+
+  it('records an authorization failure without including the query or error message', async () => {
+    create()
+    mocks.search.mockRejectedValueOnce(new OrchestrationError('forbidden', 'Private denial reason'))
+    await call('search', { query: 'private query' })
+    expect(mocks.info).toHaveBeenCalledExactlyOnceWith('Knowledge MCP tool completed', {
+      toolName: 'search',
+      operation: 'knowledge.search',
+      organizationId: 'org-1',
+      userId: 'person-1',
+      outcome: 'error',
+      durationMs: expect.any(Number),
+    })
+  })
+
+  it('distinguishes rate limiting from an executed tool', async () => {
+    create()
+    mocks.rateLimit.mockResolvedValueOnce(new Response(null, { status: 429 }))
+    await call('search', { query: 'private query' })
+    expect(mocks.search).not.toHaveBeenCalled()
+    expect(mocks.info).toHaveBeenCalledExactlyOnceWith(
+      'Knowledge MCP tool completed',
+      expect.objectContaining({ outcome: 'rate_limited' })
+    )
+    await mocks.afterResponse.mock.calls[0][0]()
+    expect(mocks.recordActivity).toHaveBeenCalledWith(
+      expect.objectContaining({
+        authKind: 'personal_api_key',
+        oauthClientId: null,
+        outcome: 'rate_limited',
+      })
+    )
+  })
+
+  it.each(['search', 'read_document', 'chat'])(
+    'records cancelled %s calls without executing the operation',
+    async (toolName) => {
+      create()
+      await call(toolName, { query: 'private query', documentId: 'doc-1' }, AbortSignal.abort())
+      expect(mocks.search).not.toHaveBeenCalled()
+      expect(mocks.read).not.toHaveBeenCalled()
+      expect(mocks.chat).not.toHaveBeenCalled()
+      expect(mocks.info).toHaveBeenCalledExactlyOnceWith(
+        'Knowledge MCP tool completed',
+        expect.objectContaining({ toolName, outcome: 'cancelled' })
+      )
+    }
+  )
+})