refactor: generic withManaged and connection tools in core

packages/core/test/unit/mcp/managed.test.ts

@@ -0,0 +1,124 @@
+import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import { describe, expect, it, vi } from 'vitest'
+import { type Backend, BackendDisposedError } from '../../../src/backend.js'
+import type { ManagedConnection, Manager } from '../../../src/manager.js'
+import { withManaged } from '../../../src/mcp/managed.js'
+
+type TestConfig = { name: string; type: string }
+type TestBackend = Backend & { ping: () => Promise<string> }
+type TestManaged = ManagedConnection<TestConfig, TestBackend>
+
+const config: TestConfig = { name: 'c', type: 'postgres' }
+
+const managed = (backend: Partial<TestBackend> = {}): TestManaged => ({
+    backend: { dispose: vi.fn(async () => {}), ping: vi.fn(async () => 'ok'), ...backend },
+    config,
+    source: 'store'
+})
+
+type FakeManager = Manager<TestConfig, TestBackend> & {
+    get: ReturnType<typeof vi.fn>
+    invalidate: ReturnType<typeof vi.fn>
+    disposeAll: ReturnType<typeof vi.fn>
+}
+
+const fakeManager = (): FakeManager => ({
+    get: vi.fn(),
+    invalidate: vi.fn(async () => {}),
+    disposeAll: vi.fn(async () => {})
+})
+
+const formatError = (cfg: TestConfig, error: unknown): string =>
+    '[' + cfg.type + '] ' + (error instanceof Error ? error.message : String(error))
+
+const text = (result: CallToolResult): string => {
+    const first = result.content[0]
+    return first && first.type === 'text' ? first.text : ''
+}
+
+describe('withManaged', () => {
+    it('returns ok with the fn result on success', async () => {
+        const manager = fakeManager()
+        manager.get.mockResolvedValue(managed())
+        const fn = vi.fn(async () => ({ value: 42 }))
+
+        const result = await withManaged(manager, 'c', fn, { formatError })
+
+        expect(result.isError).toBeFalsy()
+        expect(JSON.parse(text(result))).toEqual({ value: 42 })
+        expect(manager.get).toHaveBeenCalledTimes(1)
+    })
+
+    it('fails when manager.get throws, without invoking fn', async () => {
+        const manager = fakeManager()
+        manager.get.mockRejectedValue(new Error('no such connection'))
+        const fn = vi.fn()
+
+        const result = await withManaged(manager, 'c', fn, { formatError })
+
+        expect(result.isError).toBe(true)
+        expect(text(result)).toBe('no such connection')
+        expect(fn).not.toHaveBeenCalled()
+    })
+
+    it('retries once when the backend was disposed under the caller', async () => {
+        const manager = fakeManager()
+        manager.get.mockResolvedValue(managed())
+        const fn = vi
+            .fn()
+            .mockRejectedValueOnce(new BackendDisposedError())
+            .mockResolvedValueOnce({ value: 7 })
+
+        const result = await withManaged(manager, 'c', fn, { formatError })
+
+        expect(result.isError).toBeFalsy()
+        expect(JSON.parse(text(result))).toEqual({ value: 7 })
+        expect(manager.get).toHaveBeenCalledTimes(2)
+        expect(fn).toHaveBeenCalledTimes(2)
+    })
+
+    it('gives up with a reconfiguring message after a second stale error', async () => {
+        const manager = fakeManager()
+        manager.get.mockResolvedValue(managed())
+        const fn = vi.fn().mockRejectedValue(new BackendDisposedError())
+
+        const result = await withManaged(manager, 'c', fn, { formatError })
+
+        expect(result.isError).toBe(true)
+        expect(text(result)).toContain('being reconfigured')
+        expect(manager.get).toHaveBeenCalledTimes(2)
+        expect(fn).toHaveBeenCalledTimes(2)
+    })
+
+    it('formats non-stale errors via formatError, no retry', async () => {
+        const manager = fakeManager()
+        manager.get.mockResolvedValue(managed())
+        const fn = vi.fn().mockRejectedValue(new Error('boom'))
+
+        const result = await withManaged(manager, 'c', fn, { formatError })
+
+        expect(result.isError).toBe(true)
+        expect(text(result)).toBe('[postgres] boom')
+        expect(manager.get).toHaveBeenCalledTimes(1)
+        expect(fn).toHaveBeenCalledTimes(1)
+    })
+
+    it('honors a custom isStaleError predicate', async () => {
+        const manager = fakeManager()
+        manager.get.mockResolvedValue(managed())
+        const fn = vi
+            .fn()
+            .mockRejectedValueOnce(Object.assign(new Error('stale'), { code: 'STALE' }))
+            .mockResolvedValueOnce({ value: 1 })
+
+        const result = await withManaged(manager, 'c', fn, {
+            formatError,
+            isStaleError: (error) =>
+                error instanceof Error && (error as { code?: string }).code === 'STALE'
+        })
+
+        expect(result.isError).toBeFalsy()
+        expect(JSON.parse(text(result))).toEqual({ value: 1 })
+        expect(fn).toHaveBeenCalledTimes(2)
+    })
+})