Skip to content

Commit 7ef1087

Browse files
waleedlatif1waleedlatif1
committed
fix(search-mcp): snapshot client attribution during admission
1 parent 0f4a49e commit 7ef1087

9 files changed

Lines changed: 113 additions & 66 deletions

File tree

‎apps/sim/lib/auth/oauth-access-token.test.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ function row(overrides: Record<string, unknown> = {}) {
1818
id: 'token-1',
1919
userId: 'user-1',
2020
clientId: 'sim-cli',
21+
clientName: 'Sim CLI',
2122
scopes: ['offline_access', 'api:read'],
2223
resource: null,
2324
expiresAt: new Date(Date.now() + 60_000),
@@ -80,6 +81,7 @@ describe('verifyOAuthAccessToken', () => {
8081
kind: 'oauth_access_token',
8182
userId: 'user-1',
8283
clientId: 'sim-cli',
84+
clientName: 'Sim CLI',
8385
tokenId: 'token-1',
8486
scopes: ['offline_access', 'api:read'],
8587
expiresAt: expect.any(Date),
@@ -92,6 +94,13 @@ describe('verifyOAuthAccessToken', () => {
9294
)
9395
})
9496

97+
it('does not invent a display name for an unnamed OAuth client', async () => {
98+
queueTableRows(schemaMock.oauthAccessToken, [row({ clientName: null })])
99+
const principal = await verifyOAuthAccessToken('sim_oat_secret')
100+
expect(principal).not.toHaveProperty('clientName')
101+
expect(principal.clientId).toBe('sim-cli')
102+
})
103+
95104
it('refuses a credential that is not one of ours without a database read', async () => {
96105
expect(await reason('sim_abc')).toBe('malformed')
97106
expect(await reason('sim_oat_')).toBe('malformed')

‎apps/sim/lib/auth/oauth-access-token.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { db } from '@sim/db'
33
import { oauthAccessToken, oauthClient, user } from '@sim/db/schema'
44
import { createLogger, setRequestAuth } from '@sim/logger'
55
import { sha256Hex } from '@sim/security/hash'
6-
import { eq } from 'drizzle-orm'
6+
import { eq, sql } from 'drizzle-orm'
77
import { isAccountBlocked } from '@/lib/auth/ban'
88
import {
99
OAUTH_ACCESS_TOKEN_PREFIX,
@@ -99,6 +99,7 @@ export async function verifyOAuthAccessToken(
9999
id: oauthAccessToken.id,
100100
userId: oauthAccessToken.userId,
101101
clientId: oauthAccessToken.clientId,
102+
clientName: sql<string | null>`left(${oauthClient.name}, 256)`,
102103
scopes: oauthAccessToken.scopes,
103104
resource: oauthAccessToken.resource,
104105
expiresAt: oauthAccessToken.expiresAt,
@@ -142,6 +143,7 @@ export async function verifyOAuthAccessToken(
142143
kind: 'oauth_access_token',
143144
userId: row.userId,
144145
clientId: row.clientId,
146+
...(row.clientName ? { clientName: row.clientName } : {}),
145147
tokenId: row.id,
146148
scopes: row.scopes,
147149
expiresAt: row.expiresAt,

‎apps/sim/lib/auth/oauth-principal.test.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,4 +66,12 @@ describe('oauth_access_token principal', () => {
6666
parsePrincipal({ ...serialized, principal: { ...serialized.principal, expiresAt: 'soon' } })
6767
).toThrow('expiresAt must be an ISO timestamp')
6868
})
69+
70+
it('keeps display metadata out of persisted workflow authority', () => {
71+
const named = { ...principal, clientName: 'Registered app' }
72+
const serialized = serializePrincipal(named)
73+
expect(serialized.principal).not.toHaveProperty('clientName')
74+
expect(parsePrincipal(serialized)).toEqual(principal)
75+
expect(toPrincipalActor(named)).toEqual(toPrincipalActor(principal))
76+
})
6977
})

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

Lines changed: 59 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -525,66 +525,68 @@ describe('organization Search MCP with real ingestion and current access', () =>
525525
await db
526526
.delete(organizationSearchMcpInvocation)
527527
.where(eq(organizationSearchMcpInvocation.organizationId, organizationId))
528+
const clientName = 'MCP fixture client'.repeat(20)
528529
await db
529530
.update(oauthClient)
530-
.set({ name: 'MCP fixture client' })
531+
.set({ name: clientName })
531532
.where(eq(oauthClient.clientId, oauthClientId))
532-
await aliceOAuth.listTools()
533-
expect(fixtures.afterResponse).toHaveLength(0)
534-
await search(aliceOAuth)
535-
await value(aliceOAuth, 'read_document', { documentId })
536-
expect((await call(bob, 'read_document', { documentId })).isError).toBe(true)
537-
expect(fixtures.afterResponse).toHaveLength(3)
538-
await flushAfterResponse()
539-
const rows = await db
540-
.select()
541-
.from(organizationSearchMcpInvocation)
542-
.where(eq(organizationSearchMcpInvocation.organizationId, organizationId))
543-
.orderBy(organizationSearchMcpInvocation.createdAt)
544-
.limit(10)
545-
expect(rows).toHaveLength(3)
546-
expect(rows).toMatchObject([
547-
{
548-
organizationId,
549-
userId: aliceId,
550-
authKind: 'oauth_access_token',
551-
oauthClientId,
552-
clientName: 'MCP fixture client',
553-
toolName: 'search',
554-
outcome: 'success',
555-
},
556-
{
557-
organizationId,
558-
userId: aliceId,
559-
authKind: 'oauth_access_token',
560-
oauthClientId,
561-
clientName: 'MCP fixture client',
562-
toolName: 'read_document',
563-
outcome: 'success',
564-
},
565-
{
566-
organizationId,
567-
userId: bobId,
568-
authKind: 'personal_api_key',
569-
oauthClientId: null,
570-
clientName: null,
571-
toolName: 'read_document',
572-
outcome: 'error',
573-
},
574-
])
575-
expect(rows.every((row) => row.durationMs >= 0)).toBe(true)
576-
expect(JSON.stringify(rows)).not.toContain(documentId)
577-
expect(JSON.stringify(rows)).not.toContain(oauthTokens.alice)
578-
await db
579-
.update(oauthClient)
580-
.set({ name: 'Renamed client' })
581-
.where(eq(oauthClient.clientId, oauthClientId))
582-
const [historical] = await db
583-
.select()
584-
.from(organizationSearchMcpInvocation)
585-
.where(eq(organizationSearchMcpInvocation.id, rows[0].id))
586-
.limit(1)
587-
expect(historical.clientName).toBe('MCP fixture client')
533+
try {
534+
await aliceOAuth.listTools()
535+
expect(fixtures.afterResponse).toHaveLength(0)
536+
await search(aliceOAuth)
537+
await value(aliceOAuth, 'read_document', { documentId })
538+
expect((await call(bob, 'read_document', { documentId })).isError).toBe(true)
539+
expect(fixtures.afterResponse).toHaveLength(3)
540+
await db
541+
.update(oauthClient)
542+
.set({ name: 'Renamed client' })
543+
.where(eq(oauthClient.clientId, oauthClientId))
544+
await flushAfterResponse()
545+
const rows = await db
546+
.select()
547+
.from(organizationSearchMcpInvocation)
548+
.where(eq(organizationSearchMcpInvocation.organizationId, organizationId))
549+
.orderBy(organizationSearchMcpInvocation.createdAt)
550+
.limit(10)
551+
expect(rows).toHaveLength(3)
552+
expect(rows).toMatchObject([
553+
{
554+
organizationId,
555+
userId: aliceId,
556+
authKind: 'oauth_access_token',
557+
oauthClientId,
558+
clientName: clientName.slice(0, 256),
559+
toolName: 'search',
560+
outcome: 'success',
561+
},
562+
{
563+
organizationId,
564+
userId: aliceId,
565+
authKind: 'oauth_access_token',
566+
oauthClientId,
567+
clientName: clientName.slice(0, 256),
568+
toolName: 'read_document',
569+
outcome: 'success',
570+
},
571+
{
572+
organizationId,
573+
userId: bobId,
574+
authKind: 'personal_api_key',
575+
oauthClientId: null,
576+
clientName: null,
577+
toolName: 'read_document',
578+
outcome: 'error',
579+
},
580+
])
581+
expect(rows.every((row) => row.durationMs >= 0)).toBe(true)
582+
expect(JSON.stringify(rows)).not.toContain(documentId)
583+
expect(JSON.stringify(rows)).not.toContain(oauthTokens.alice)
584+
} finally {
585+
await db
586+
.update(oauthClient)
587+
.set({ name: 'Search MCP OAuth fixture' })
588+
.where(eq(oauthClient.clientId, oauthClientId))
589+
}
588590
})
589591

590592
it('enforces current document and organization access on Search OAuth clients', async () => {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ const activity: SearchMcpActivityInput = {
1919
userId: 'actor',
2020
authKind: 'personal_api_key',
2121
oauthClientId: null,
22+
clientName: null,
2223
toolName: 'read_document',
2324
outcome: 'success',
2425
durationMs: 42,

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { db } from '@sim/db'
2-
import { oauthClient, organizationSearchMcpInvocation } from '@sim/db/schema'
2+
import { organizationSearchMcpInvocation } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { getErrorMessage } from '@sim/utils/errors'
55
import { generateId } from '@sim/utils/id'
@@ -13,6 +13,7 @@ export type SearchMcpActivityInput = Pick<
1313
| 'userId'
1414
| 'authKind'
1515
| 'oauthClientId'
16+
| 'clientName'
1617
| 'toolName'
1718
| 'outcome'
1819
| 'durationMs'
@@ -32,9 +33,7 @@ export async function recordOrganizationSearchMcpActivity(
3233
userId: input.userId,
3334
authKind: input.authKind,
3435
oauthClientId: input.oauthClientId,
35-
clientName: input.oauthClientId
36-
? sql`(SELECT left(${oauthClient.name}, 256) FROM ${oauthClient} WHERE ${oauthClient.clientId} = ${input.oauthClientId})`
37-
: null,
36+
clientName: input.clientName,
3837
toolName: input.toolName,
3938
outcome: input.outcome,
4039
durationMs: input.durationMs,

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ function payload(result: CallToolResult): unknown {
8282
beforeEach(() => {
8383
vi.clearAllMocks()
8484
mocks.tools.clear()
85-
mocks.rateLimit.mockResolvedValue(null)
85+
mocks.rateLimit.mockReset().mockResolvedValue(null)
8686
mocks.search.mockResolvedValue({ results: [] })
8787
mocks.read.mockResolvedValue({
8888
knowledgeBaseId: 'index-1',
@@ -339,6 +339,7 @@ describe('MCP tool completion records', () => {
339339
kind: 'oauth_access_token',
340340
userId: 'oauth-person',
341341
clientId: 'registered-client',
342+
clientName: 'Registered app',
342343
tokenId: 'private-token-id',
343344
scopes: ['search:read'],
344345
expiresAt: new Date(Date.now() + 60000),
@@ -355,6 +356,7 @@ describe('MCP tool completion records', () => {
355356
userId: 'oauth-person',
356357
authKind: 'oauth_access_token',
357358
oauthClientId: 'registered-client',
359+
clientName: 'Registered app',
358360
toolName: 'search',
359361
outcome: 'success',
360362
durationMs: expect.any(Number),
@@ -426,7 +428,9 @@ describe('MCP tool completion records', () => {
426428
'records cancelled %s calls without executing the operation',
427429
async (toolName) => {
428430
create()
431+
mocks.rateLimit.mockResolvedValueOnce(new Response(null, { status: 429 }))
429432
await call(toolName, { query: 'private query', documentId: 'doc-1' }, AbortSignal.abort())
433+
expect(mocks.rateLimit).not.toHaveBeenCalled()
430434
expect(mocks.search).not.toHaveBeenCalled()
431435
expect(mocks.read).not.toHaveBeenCalled()
432436
expect(mocks.chat).not.toHaveBeenCalled()
@@ -436,4 +440,19 @@ describe('MCP tool completion records', () => {
436440
)
437441
}
438442
)
443+
444+
it('records cancellation during rate-limit admission instead of an exhausted bucket', async () => {
445+
create()
446+
const controller = new AbortController()
447+
mocks.rateLimit.mockImplementationOnce(async () => {
448+
controller.abort()
449+
return new Response(null, { status: 429 })
450+
})
451+
await call('search', { query: 'private query' }, controller.signal)
452+
expect(mocks.search).not.toHaveBeenCalled()
453+
await mocks.afterResponse.mock.calls[0][0]()
454+
expect(mocks.recordActivity).toHaveBeenCalledWith(
455+
expect.objectContaining({ outcome: 'cancelled' })
456+
)
457+
})
439458
})

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,9 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
8383
const signal = AbortSignal.any([request.signal, toolSignal])
8484
let outcome: SearchMcpActivityInput['outcome'] = 'error'
8585
try {
86+
signal.throwIfAborted()
8687
const limited = await v2RateLimits.publicApi.enforce(request, auth, operation)
88+
signal.throwIfAborted()
8789
if (limited) {
8890
outcome = 'rate_limited'
8991
const retryAfter = parseRetryAfter(
@@ -96,7 +98,6 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
9698
: `API rate limit exceeded. Retry in ${Math.ceil(retryAfter / 1000)} seconds.`
9799
)
98100
}
99-
signal.throwIfAborted()
100101
const result = await run(new ResolvedSecretTraceRegistry(), signal)
101102
outcome = signal.aborted ? 'cancelled' : result.isError ? 'error' : 'success'
102103
return result
@@ -122,6 +123,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
122123
userId: resolvePrincipalSubjectUserId(principal) ?? null,
123124
authKind: principal.kind,
124125
oauthClientId: principal.kind === 'oauth_access_token' ? principal.clientId : null,
126+
clientName: principal.kind === 'oauth_access_token' ? (principal.clientName ?? null) : null,
125127
outcome,
126128
durationMs: Math.round(performance.now() - startedAt),
127129
createdAt: new Date(),

‎packages/auth/src/principal.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ export interface OAuthAccessTokenPrincipal {
5252
kind: 'oauth_access_token'
5353
userId: string
5454
clientId: string
55+
/** Admission-time display metadata only; never grants authority or enters workflow payloads. */
56+
clientName?: string
5557
/** The `oauth_access_token` row id, never the token itself. */
5658
tokenId: string
5759
scopes: readonly string[]
@@ -291,7 +293,7 @@ export type WorkflowExecutionPrincipal =
291293
type SerializedWorkflowExecutionPrincipal =
292294
| SessionPrincipal
293295
| PersonalApiKeyPrincipal
294-
| (Omit<OAuthAccessTokenPrincipal, 'expiresAt'> & { expiresAt: string })
296+
| (Omit<OAuthAccessTokenPrincipal, 'expiresAt' | 'clientName'> & { expiresAt: string })
295297
| WorkspaceApiKeyPrincipal
296298
| SystemPrincipal
297299
| (Omit<SubjectDelegatedPrincipal, 'issuedAt' | 'expiresAt'> & {
@@ -398,7 +400,10 @@ export function serializePrincipal(principal: WorkflowExecutionPrincipal): Seria
398400
return {
399401
version: 1,
400402
principal: {
401-
...principal,
403+
kind: principal.kind,
404+
userId: principal.userId,
405+
clientId: principal.clientId,
406+
tokenId: principal.tokenId,
402407
scopes: [...principal.scopes],
403408
expiresAt: principal.expiresAt.toISOString(),
404409
},

0 commit comments

Comments
 (0)