Address PR review feedback (#5754)

- Make cache invalidation retries ownership-safe
- Preserve proven newer successes for superseded refreshes

Author: j15z

apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts

@@ -333,6 +333,38 @@ describe('MCP server refresh route', () => {
     expect(mockUpdate).not.toHaveBeenCalled()
   })
 
+  it('preserves a newer successful refresh when discovery is superseded', async () => {
+    mockDiscoverServerTools.mockResolvedValueOnce({
+      tools: [],
+      state: 'superseded',
+    })
+    const newerSuccessfulServer = {
+      ...initialServer,
+      lastConnected: new Date(initialServer.lastConnected.getTime() + 60_000),
+      lastToolsRefresh: new Date(initialServer.lastToolsRefresh.getTime() + 60_000),
+      toolCount: 7,
+    }
+    mockSelect.mockReturnValueOnce(selectRows([newerSuccessfulServer]))
+
+    const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', {
+      method: 'POST',
+    }) as NextRequest
+    const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) })
+    const body = await response.json()
+
+    expect(body.data).toEqual(
+      expect.objectContaining({
+        status: 'connected',
+        error: null,
+        toolCount: 7,
+        workflowsUpdated: 0,
+        updatedWorkflowIds: [],
+      })
+    )
+    expect(mockSelect).toHaveBeenCalledTimes(2)
+    expect(mockUpdate).not.toHaveBeenCalled()
+  })
+
   it('fails closed without syncing workflows when discovery is superseded', async () => {
     mockDiscoverServerTools.mockResolvedValueOnce({
       tools: [],

apps/sim/app/api/mcp/servers/[id]/refresh/route.ts

@@ -265,8 +265,15 @@ export const POST = withRouteHandler(
         let connectionStatus = refreshedServer?.connectionStatus ?? 'error'
         let lastError = refreshedServer ? refreshedServer.lastError : discoveryError
         let toolCount = refreshedServer?.toolCount ?? discoveredTools.length
-
-        if (discoveryState === 'superseded') {
+        const newerSuccessWonRace =
+          connectionStatus === 'connected' &&
+          refreshedServer?.lastToolsRefresh != null &&
+          (server.lastToolsRefresh == null ||
+            refreshedServer.lastToolsRefresh > server.lastToolsRefresh) &&
+          refreshedServer.lastConnected != null &&
+          (server.lastConnected == null || refreshedServer.lastConnected > server.lastConnected)
+
+        if (discoveryState === 'superseded' && !newerSuccessWonRace) {
           connectionStatus = 'disconnected'
           lastError = 'Tool discovery was superseded by a newer refresh. Please retry.'
           toolCount = 0
@@ -280,13 +287,6 @@ export const POST = withRouteHandler(
         }
 
         if (discoveryError !== null && connectionStatus === 'connected') {
-          const newerSuccessWonRace =
-            refreshedServer?.lastToolsRefresh != null &&
-            (server.lastToolsRefresh == null ||
-              refreshedServer.lastToolsRefresh > server.lastToolsRefresh) &&
-            refreshedServer.lastConnected != null &&
-            (server.lastConnected == null || refreshedServer.lastConnected > server.lastConnected)
-
           if (!newerSuccessWonRace) {
             connectionStatus = 'disconnected'
             lastError = oauthAuthorizationRequired ? null : discoveryError

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

@@ -716,7 +716,7 @@ describe('McpService.discoverTools per-server caching', () => {
     expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential)
   })
 
-  it('best-effort deletes both cache keys when invalidation cannot be ordered', async () => {
+  it('preserves cache and database state when invalidation cannot be ordered', async () => {
     const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a`
     const failureKey = `${serverKey}:failure`
     mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
@@ -725,19 +725,21 @@ describe('McpService.discoverTools per-server caching', () => {
       expiry: Date.now() + 60_000,
     })
     cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 })
-    mockCacheAdapter.beginMutation
-      .mockRejectedValueOnce(new Error('cache ordering unavailable'))
-      .mockRejectedValueOnce(new Error('cache ordering unavailable'))
+    for (let attempt = 0; attempt < 6; attempt++) {
+      mockCacheAdapter.beginMutation.mockRejectedValueOnce(new Error('cache ordering unavailable'))
+    }
 
     await mcpService.clearCache(WORKSPACE_ID)
 
-    expect(cacheStore.has(serverKey)).toBe(false)
-    expect(cacheStore.has(failureKey)).toBe(false)
-    expect(mockCacheAdapter.delete).toHaveBeenCalledWith(serverKey)
-    expect(mockCacheAdapter.delete).toHaveBeenCalledWith(failureKey)
+    expect(cacheStore.get(serverKey)?.tools).toEqual([tool('stale-tool', 'mcp-a')])
+    expect(cacheStore.has(failureKey)).toBe(true)
+    expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(6)
+    expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalled()
+    expect(mockCacheAdapter.delete).not.toHaveBeenCalled()
+    expect(mockUpdateSet).not.toHaveBeenCalled()
   })
 
-  it('best-effort deletes both cache keys and publishes the barrier when ordered invalidation fails', async () => {
+  it('reacquires ownership after a transient atomic invalidation failure', async () => {
     const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a`
     const failureKey = `${serverKey}:failure`
     mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
@@ -752,15 +754,77 @@ describe('McpService.discoverTools per-server caching', () => {
 
     await mcpService.clearCache(WORKSPACE_ID)
 
-    const mutationId = mockCacheAdapter.applyMutationIfCurrent.mock.calls[0][1]
+    const firstMutationId = mockCacheAdapter.applyMutationIfCurrent.mock.calls[0][1]
+    const successfulMutationId = mockCacheAdapter.applyMutationIfCurrent.mock.calls[1][1]
+    expect(successfulMutationId).toBeGreaterThan(firstMutationId)
+    expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(2)
+    expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(2)
     expect(cacheStore.has(serverKey)).toBe(false)
     expect(cacheStore.has(failureKey)).toBe(false)
-    expect(mockCacheAdapter.delete).toHaveBeenCalledWith(serverKey)
-    expect(mockCacheAdapter.delete).toHaveBeenCalledWith(failureKey)
+    expect(mockCacheAdapter.delete).not.toHaveBeenCalled()
     expect(mockUpdateSet).toHaveBeenCalledWith({
       toolCount: 0,
-      lastToolsRefresh: new Date(mutationId),
+      lastToolsRefresh: new Date(successfulMutationId),
+    })
+  })
+
+  it('preserves a newer winner acquired between invalidation retry begin and apply', async () => {
+    const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a`
+    const failureKey = `${serverKey}:failure`
+    const winner = tool('winner-tool', 'mcp-a')
+    mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
+    cacheStore.set(serverKey, {
+      tools: [tool('stale-tool', 'mcp-a')],
+      expiry: Date.now() + 60_000,
     })
+    cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 })
+    const defaultApply = mockCacheAdapter.applyMutationIfCurrent.getMockImplementation()
+    mockCacheAdapter.applyMutationIfCurrent
+      .mockRejectedValueOnce(new Error('atomic invalidation unavailable'))
+      .mockImplementationOnce(async (scopeKey, mutationId, setEntry, deleteKeys) => {
+        const newerMutationId = await mockCacheAdapter.beginMutation(scopeKey)
+        await defaultApply?.(
+          scopeKey,
+          newerMutationId,
+          { key: serverKey, tools: [winner], ttlMs: 60_000 },
+          [failureKey]
+        )
+        return defaultApply?.(scopeKey, mutationId, setEntry, deleteKeys) ?? false
+      })
+
+    await mcpService.clearCache(WORKSPACE_ID)
+
+    expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(2)
+    expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3)
+    expect(cacheStore.get(serverKey)?.tools).toEqual([winner])
+    expect(cacheStore.has(failureKey)).toBe(false)
+    expect(mockCacheAdapter.delete).not.toHaveBeenCalled()
+    expect(mockUpdateSet).not.toHaveBeenCalled()
+  })
+
+  it('bounds persistent atomic invalidation failures without raw deletes', async () => {
+    const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a`
+    const failureKey = `${serverKey}:failure`
+    const reflectedCredential = 'opaque-cache-provider-message'
+    mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
+    cacheStore.set(serverKey, {
+      tools: [tool('stale-tool', 'mcp-a')],
+      expiry: Date.now() + 60_000,
+    })
+    cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 })
+    for (let attempt = 0; attempt < 3; attempt++) {
+      mockCacheAdapter.applyMutationIfCurrent.mockRejectedValueOnce(new Error(reflectedCredential))
+    }
+
+    await mcpService.clearCache(WORKSPACE_ID)
+
+    expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3)
+    expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(3)
+    expect(cacheStore.get(serverKey)?.tools).toEqual([tool('stale-tool', 'mcp-a')])
+    expect(cacheStore.has(failureKey)).toBe(true)
+    expect(mockCacheAdapter.delete).not.toHaveBeenCalled()
+    expect(mockUpdateSet).not.toHaveBeenCalled()
+    expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential)
   })
 
   it('preserves a newer cache winner when invalidation is superseded', async () => {

apps/sim/lib/mcp/service.ts

@@ -68,6 +68,7 @@ export function getTimestampMillisecondBounds(timestamp: string): {
 
 const FAILURE_CACHE_SENTINEL: McpTool[] = []
 const CACHE_MUTATION_BEGIN_ATTEMPTS = 2
+const CACHE_INVALIDATION_ATTEMPTS = 3
 const STATUS_UPDATE_CAS_ATTEMPTS = 3
 const STATUS_METADATA_RACE_RETRY_ATTEMPTS = 3
 
@@ -702,45 +703,35 @@ class McpService {
     return null
   }
 
-  private async bestEffortInvalidateServerCache(
-    workspaceId: string,
-    serverId: string
-  ): Promise<void> {
+  private async invalidateServerCache(workspaceId: string, serverId: string): Promise<void> {
     const keys = [serverCacheKey(workspaceId, serverId), failureCacheKey(workspaceId, serverId)]
-    const results = await Promise.allSettled(keys.map((key) => this.cacheAdapter.delete(key)))
-    const failures = results.flatMap((result) =>
-      result.status === 'rejected' ? [getMcpSafeErrorDiagnostics(result.reason)] : []
-    )
-    if (failures.length > 0) {
-      logger.warn(`Failed best-effort cache invalidation for server ${serverId}`, {
+    for (let attempt = 0; attempt < CACHE_INVALIDATION_ATTEMPTS; attempt++) {
+      // Reacquire a fresh ownership token after every unavailable transition.
+      // This orders the retry after work that completed during the unknown
+      // attempt, while work starting later can still supersede the new token.
+      const mutation = await this.beginServerCacheMutation(workspaceId, serverId)
+      if (!mutation) continue
+
+      const cacheApplied = await this.applyServerCacheMutation(
         workspaceId,
-        failures,
-      })
-    }
-  }
+        serverId,
+        mutation,
+        null,
+        keys
+      )
+      if (cacheApplied === 'superseded') return
+      if (cacheApplied === 'unavailable') continue
 
-  private async invalidateServerCache(workspaceId: string, serverId: string): Promise<void> {
-    const mutation = await this.beginServerCacheMutation(workspaceId, serverId)
-    if (!mutation) {
-      // During a total cache-backend outage there is no cross-process token we
-      // can advance, so this can only be best effort. Deleting both keys still
-      // prevents stale reads whenever the backend accepts ordinary commands.
-      await this.bestEffortInvalidateServerCache(workspaceId, serverId)
+      await this.markServerCacheInvalidated(serverId, workspaceId, new Date(mutation.id))
       return
     }
-    const cacheApplied = await this.applyServerCacheMutation(
-      workspaceId,
-      serverId,
-      mutation,
-      null,
-      [serverCacheKey(workspaceId, serverId), failureCacheKey(workspaceId, serverId)]
-    )
-    if (cacheApplied === 'superseded') return
-    if (cacheApplied === 'unavailable') {
-      await this.bestEffortInvalidateServerCache(workspaceId, serverId)
-    }
 
-    await this.markServerCacheInvalidated(serverId, workspaceId, new Date(mutation.id))
+    // Without an ownership-checked transition, deleting either key can erase
+    // a newer process's winner. Leave cache and database state unchanged.
+    logger.warn(`Cache invalidation unavailable for server ${serverId}`, {
+      workspaceId,
+      attempts: CACHE_INVALIDATION_ATTEMPTS,
+    })
   }
 
   private async getCurrentCachedTools(