Skip to content

Commit 6ff6634

Browse files
j15zj15z
committed
Unify MCP cache and database publication order (#5754)
Use timestamp-based mutation tokens for both cache ownership and database CAS ordering, and invalidate Redis mutation owners before deleting entries during a full clear.
1 parent 80da6da commit 6ff6634

7 files changed

Lines changed: 167 additions & 78 deletions

File tree

‎apps/sim/lib/mcp/service.test.ts‎

Lines changed: 62 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,8 @@ const {
4848
cacheStore.delete(key)
4949
}),
5050
beginMutation: vi.fn(async (scopeKey: string) => {
51-
const mutationId = ++nextMutationId
51+
const mutationId = Math.max(nextMutationId + 1, Date.now())
52+
nextMutationId = mutationId
5253
cacheMutations.set(scopeKey, mutationId)
5354
return mutationId
5455
}),
@@ -89,10 +90,12 @@ const {
8990
}
9091
),
9192
clear: vi.fn(async () => {
92-
cacheStore.clear()
9393
for (const scopeKey of cacheMutations.keys()) {
94-
cacheMutations.set(scopeKey, ++nextMutationId)
94+
const mutationId = Math.max(nextMutationId + 1, Date.now())
95+
nextMutationId = mutationId
96+
cacheMutations.set(scopeKey, mutationId)
9597
}
98+
cacheStore.clear()
9699
}),
97100
dispose: () => {},
98101
}
@@ -560,6 +563,56 @@ describe('McpService.discoverTools per-server caching', () => {
560563
expect(mockUpdateSet).toHaveBeenCalledTimes(1)
561564
})
562565

566+
it('uses the cache mutation token to order database publication after a begin retry', async () => {
567+
vi.useFakeTimers({ toFake: ['Date'] })
568+
try {
569+
mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
570+
571+
let rejectOlderBegin: ((error: Error) => void) | undefined
572+
mockCacheAdapter.beginMutation.mockImplementationOnce(
573+
() =>
574+
new Promise<number>((_resolve, reject) => {
575+
rejectOlderBegin = reject
576+
})
577+
)
578+
mockListTools
579+
.mockRejectedValueOnce(new Error('Later-started discovery failed'))
580+
.mockResolvedValueOnce([tool('retry-winner', 'mcp-a')])
581+
582+
vi.setSystemTime(new Date('2030-02-01T00:00:00.000Z'))
583+
const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false)
584+
await vi.waitFor(() => expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(1))
585+
586+
const laterMutationTime = new Date('2030-02-01T00:00:01.000Z')
587+
vi.setSystemTime(laterMutationTime)
588+
const later = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true)
589+
await expect(later).rejects.toThrow('Later-started discovery failed')
590+
591+
const retriedMutationTime = new Date('2030-02-01T00:00:02.000Z')
592+
vi.setSystemTime(retriedMutationTime)
593+
rejectOlderBegin?.(new Error('Transient mutation start failure'))
594+
await expect(older).resolves.toEqual([tool('retry-winner', 'mcp-a')])
595+
596+
const publications = mockUpdateSet.mock.calls
597+
.map(([update]) => update)
598+
.filter((update) => update.lastToolsRefresh)
599+
expect(publications.map((update) => update.connectionStatus)).toEqual([
600+
'disconnected',
601+
'connected',
602+
])
603+
expect(publications.map((update) => update.lastToolsRefresh)).toEqual([
604+
laterMutationTime,
605+
retriedMutationTime,
606+
])
607+
expect(cacheStore.get(`workspace:${WORKSPACE_ID}:server:mcp-a`)?.tools).toEqual([
608+
tool('retry-winner', 'mcp-a'),
609+
])
610+
expect(cacheStore.has(`workspace:${WORKSPACE_ID}:server:mcp-a:failure`)).toBe(false)
611+
} finally {
612+
vi.useRealTimers()
613+
}
614+
})
615+
563616
it('fails a newer discovery closed while an older mutation remains the owner', async () => {
564617
const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a`
565618
mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
@@ -763,29 +816,30 @@ describe('McpService.discoverTools per-server caching', () => {
763816
})
764817
)
765818

766-
const olderStartedAt = new Date('2026-02-01T00:00:00.000Z')
819+
const olderStartedAt = new Date('2030-02-01T00:00:00.000Z')
767820
vi.setSystemTime(olderStartedAt)
768821
const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false)
769822
await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1))
770823

771-
const newerStartedAt = new Date('2026-02-01T00:00:01.000Z')
824+
const newerStartedAt = new Date('2030-02-01T00:00:01.000Z')
772825
vi.setSystemTime(newerStartedAt)
773826
const newer = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true)
774827
await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(2))
775828

776-
vi.setSystemTime(new Date('2026-02-01T00:00:02.000Z'))
829+
vi.setSystemTime(new Date('2030-02-01T00:00:02.000Z'))
777830
resolveOlder?.([tool('old-tool', 'mcp-a')])
778831
await expect(older).resolves.toEqual([])
779832

780-
vi.setSystemTime(new Date('2026-02-01T00:00:03.000Z'))
833+
vi.setSystemTime(new Date('2030-02-01T00:00:03.000Z'))
781834
resolveNewer?.([tool('new-tool', 'mcp-a')])
782835
await expect(newer).resolves.toEqual([tool('new-tool', 'mcp-a')])
783836

784837
const publishedRefreshTimes = mockUpdateSet.mock.calls
785838
.map(([update]) => update)
786839
.filter((update) => update.connectionStatus === 'connected')
787840
.map((update) => update.lastToolsRefresh)
788-
expect(publishedRefreshTimes).toEqual([newerStartedAt])
841+
expect(publishedRefreshTimes).toHaveLength(1)
842+
expect(publishedRefreshTimes[0].getTime()).toBeGreaterThanOrEqual(newerStartedAt.getTime())
789843
} finally {
790844
vi.useRealTimers()
791845
}

‎apps/sim/lib/mcp/service.ts‎

Lines changed: 26 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -99,13 +99,13 @@ type ServerStatusUpdate =
9999
outcome: 'connected'
100100
toolCount: number
101101
configUpdatedAt: string
102-
discoveryStartedAt: Date
102+
publicationOrder: Date
103103
}
104104
| {
105105
outcome: 'failed'
106106
error: string
107107
configUpdatedAt: string
108-
discoveryStartedAt: Date
108+
publicationOrder: Date
109109
}
110110

111111
function isOauthAuthorizationError(error: unknown, authType: McpServerConfig['authType']): boolean {
@@ -429,7 +429,7 @@ class McpService {
429429
lt(mcpServers.updatedAt, configUpdatedAt.endExclusive),
430430
or(
431431
isNull(mcpServers.lastToolsRefresh),
432-
lte(mcpServers.lastToolsRefresh, update.discoveryStartedAt)
432+
lte(mcpServers.lastToolsRefresh, update.publicationOrder)
433433
)
434434
)
435435

@@ -441,10 +441,10 @@ class McpService {
441441
lastConnected: now,
442442
lastError: null,
443443
toolCount: update.toolCount,
444-
// This column is also the publication ordering token. Persist the
445-
// discovery start (rather than finish) so a newer-started run can
446-
// still publish after an older run completes while it is in flight.
447-
lastToolsRefresh: update.discoveryStartedAt,
444+
// The cache mutation id is a monotonic millisecond timestamp. Use
445+
// that same token here so cache and database publication cannot
446+
// choose different winners during begin-mutation retries.
447+
lastToolsRefresh: update.publicationOrder,
448448
statusConfig: {
449449
consecutiveFailures: 0,
450450
lastSuccessfulDiscovery: now.toISOString(),
@@ -485,7 +485,7 @@ class McpService {
485485
connectionStatus: isErrorState ? 'error' : 'disconnected',
486486
lastError: update.error || 'Unknown error',
487487
toolCount: 0,
488-
lastToolsRefresh: update.discoveryStartedAt,
488+
lastToolsRefresh: update.publicationOrder,
489489
statusConfig: {
490490
consecutiveFailures: newFailures,
491491
lastSuccessfulDiscovery: currentConfig.lastSuccessfulDiscovery,
@@ -538,7 +538,7 @@ class McpService {
538538
serverId: string,
539539
workspaceId: string,
540540
configUpdatedAt: string,
541-
discoveryStartedAt: Date
541+
publicationOrder: Date
542542
): Promise<boolean> {
543543
try {
544544
const configUpdatedAtBounds = getTimestampMillisecondBounds(configUpdatedAt)
@@ -548,7 +548,7 @@ class McpService {
548548
connectionStatus: 'disconnected',
549549
lastError: null,
550550
toolCount: 0,
551-
lastToolsRefresh: discoveryStartedAt,
551+
lastToolsRefresh: publicationOrder,
552552
})
553553
.where(
554554
and(
@@ -559,7 +559,7 @@ class McpService {
559559
lt(mcpServers.updatedAt, configUpdatedAtBounds.endExclusive),
560560
or(
561561
isNull(mcpServers.lastToolsRefresh),
562-
lte(mcpServers.lastToolsRefresh, discoveryStartedAt)
562+
lte(mcpServers.lastToolsRefresh, publicationOrder)
563563
)
564564
)
565565
)
@@ -636,8 +636,7 @@ class McpService {
636636
workspaceId: string,
637637
config: McpServerConfig,
638638
mutation: CacheMutation | null,
639-
tools: McpTool[],
640-
discoveryStartedAt: Date
639+
tools: McpTool[]
641640
): Promise<boolean> {
642641
const cacheApplied = await this.applyServerCacheMutation(
643642
workspaceId,
@@ -650,13 +649,13 @@ class McpService {
650649
},
651650
[failureCacheKey(workspaceId, config.id)]
652651
)
653-
if (!cacheApplied) return false
652+
if (!cacheApplied || !mutation) return false
654653

655654
const statusApplied = await this.updateServerStatus(config.id, workspaceId, {
656655
outcome: 'connected',
657656
toolCount: tools.length,
658657
configUpdatedAt: config.updatedAt!,
659-
discoveryStartedAt,
658+
publicationOrder: new Date(mutation.id),
660659
})
661660
if (statusApplied) return true
662661

@@ -675,8 +674,7 @@ class McpService {
675674
config: McpServerConfig,
676675
mutation: CacheMutation | null,
677676
error: unknown,
678-
message: string,
679-
discoveryStartedAt: Date
677+
message: string
680678
): Promise<boolean> {
681679
const cacheApplied = await this.applyServerCacheMutation(
682680
workspaceId,
@@ -691,13 +689,13 @@ class McpService {
691689
},
692690
[serverCacheKey(workspaceId, config.id)]
693691
)
694-
if (!cacheApplied) return false
692+
if (!cacheApplied || !mutation) return false
695693

696694
const statusApplied = await this.updateServerStatus(config.id, workspaceId, {
697695
outcome: 'failed',
698696
error: message,
699697
configUpdatedAt: config.updatedAt!,
700-
discoveryStartedAt,
698+
publicationOrder: new Date(mutation.id),
701699
})
702700
if (statusApplied) return true
703701

@@ -712,8 +710,7 @@ class McpService {
712710
private async publishOauthPending(
713711
workspaceId: string,
714712
config: McpServerConfig,
715-
mutation: CacheMutation | null,
716-
discoveryStartedAt: Date
713+
mutation: CacheMutation | null
717714
): Promise<boolean> {
718715
const cacheApplied = await this.applyServerCacheMutation(
719716
workspaceId,
@@ -722,13 +719,13 @@ class McpService {
722719
null,
723720
[serverCacheKey(workspaceId, config.id), failureCacheKey(workspaceId, config.id)]
724721
)
725-
if (!cacheApplied) return false
722+
if (!cacheApplied || !mutation) return false
726723

727724
return this.markServerOauthPending(
728725
config.id,
729726
workspaceId,
730727
config.updatedAt!,
731-
discoveryStartedAt
728+
new Date(mutation.id)
732729
)
733730
}
734731

@@ -738,8 +735,6 @@ class McpService {
738735
forceRefresh = false
739736
): Promise<McpTool[]> {
740737
const requestId = generateRequestId()
741-
const discoveryStartedAt = new Date()
742-
743738
try {
744739
logger.info(`[${requestId}] Discovering MCP tools for workspace ${workspaceId}`)
745740

@@ -822,8 +817,7 @@ class McpService {
822817
workspaceId,
823818
outcome.resolvedConfig,
824819
outcome.mutation,
825-
outcome.tools,
826-
discoveryStartedAt
820+
outcome.tools
827821
)
828822
if (!published) {
829823
logger.info(
@@ -847,12 +841,7 @@ class McpService {
847841
logger.info(
848842
`[${requestId}] Skipping server ${server.name}: OAuth authorization pending`
849843
)
850-
await this.publishOauthPending(
851-
workspaceId,
852-
outcome.config,
853-
outcome.mutation,
854-
discoveryStartedAt
855-
)
844+
await this.publishOauthPending(workspaceId, outcome.config, outcome.mutation)
856845
return { tools: [], cached: 0, fetched: 0, failed: 0, liveConnection: null }
857846
}
858847
if (outcome.kind === 'unhealthy') {
@@ -867,8 +856,7 @@ class McpService {
867856
outcome.config,
868857
outcome.mutation,
869858
outcome.originalError,
870-
outcome.message,
871-
discoveryStartedAt
859+
outcome.message
872860
)
873861
return { tools: [], cached: 0, fetched: 0, failed: 1, liveConnection: null }
874862
})
@@ -941,7 +929,6 @@ class McpService {
941929
forceRefresh: boolean
942930
): Promise<McpTool[]> {
943931
const requestId = generateRequestId()
944-
const discoveryStartedAt = new Date()
945932
const maxRetries = 2
946933

947934
if (!forceRefresh) {
@@ -991,8 +978,7 @@ class McpService {
991978
workspaceId,
992979
resolvedConfig,
993980
mutation,
994-
tools,
995-
discoveryStartedAt
981+
tools
996982
)
997983
if (!published) {
998984
logger.info(
@@ -1015,15 +1001,14 @@ class McpService {
10151001
}
10161002
if (config) {
10171003
if (isOauthAuthorizationError(error, config.authType)) {
1018-
await this.publishOauthPending(workspaceId, config, mutation, discoveryStartedAt)
1004+
await this.publishOauthPending(workspaceId, config, mutation)
10191005
} else {
10201006
await this.publishFailedDiscovery(
10211007
workspaceId,
10221008
config,
10231009
mutation,
10241010
error,
1025-
getDiscoveryFailureMessage(error, config.authType, 'Connection failed'),
1026-
discoveryStartedAt
1011+
getDiscoveryFailureMessage(error, config.authType, 'Connection failed')
10271012
)
10281013
}
10291014
}

‎apps/sim/lib/mcp/storage/adapter.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,10 @@ export interface McpCacheStorageAdapter {
1616
set(key: string, tools: McpTool[], ttlMs: number): Promise<void>
1717
delete(key: string): Promise<void>
1818
/**
19-
* Starts an ordered mutation for one server. Conditional writes using an
20-
* older mutation id must be ignored so a slow discovery cannot overwrite a
21-
* newer result.
19+
* Starts an ordered mutation for one server and returns a monotonic Unix
20+
* timestamp in milliseconds. Conditional writes using an older mutation id
21+
* must be ignored so a slow discovery cannot overwrite a newer result. The
22+
* same value orders database publication for an end-to-end consistent state.
2223
*/
2324
beginMutation(scopeKey: string): Promise<number>
2425
setIfCurrentMutation(

‎apps/sim/lib/mcp/storage/memory-cache.test.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,9 +167,13 @@ describe('MemoryMcpCache', () => {
167167

168168
describe('ordered mutations', () => {
169169
it('prevents an older discovery from overwriting a newer cache result', async () => {
170+
const beforeMutation = Date.now()
170171
const older = await cache.beginMutation('server-1')
171172
const newer = await cache.beginMutation('server-1')
172173

174+
expect(older).toBeGreaterThanOrEqual(beforeMutation)
175+
expect(newer).toBeGreaterThan(older)
176+
173177
expect(
174178
await cache.setIfCurrentMutation(
175179
'server-1',

0 commit comments

Comments
 (0)