refactor: generic withManaged and connection tools in core

packages/core/src/index.ts

@@ -22,6 +22,10 @@ export {
 } from './format.js'
 export type { ManagedConnection, Manager, ManagerDeps } from './manager.js'
 export { createManager } from './manager.js'
+export type { ConnectionToolsDeps } from './mcp/connectionTools.js'
+export { registerConnectionTools } from './mcp/connectionTools.js'
+export type { WithManagedDeps } from './mcp/managed.js'
+export { withManaged } from './mcp/managed.js'
 export { errorMessage, fail, ok } from './mcp/respond.js'
 export type { Tunnel } from './net/tunnel.js'
 export { buildSshAuth, expandHome, openTunnel } from './net/tunnel.js'

packages/core/src/mcp/connectionTools.ts

@@ -0,0 +1,117 @@
+import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
+import * as z from 'zod'
+import type { Backend } from '../backend.js'
+import { BackendDisposedError } from '../backend.js'
+import type { ResolvedConnection } from '../config/baseSchema.js'
+import type { Registry } from '../config/registry.js'
+import type { ManagedConnection, Manager } from '../manager.js'
+import { errorMessage, fail, ok } from './respond.js'
+
+export type ConnectionToolsDeps<TConfig extends { name: string }, TBackend extends Backend> = {
+    registry: Registry<TConfig>
+    manager: Manager<TConfig, TBackend>
+    fullSchema: z.ZodType<TConfig, z.ZodTypeDef, unknown>
+    patchSchema: z.ZodTypeAny
+    publicView: (resolved: ResolvedConnection<TConfig>) => unknown
+    ping: (managed: ManagedConnection<TConfig, TBackend>) => Promise<Record<string, unknown>>
+    formatError: (config: TConfig, error: unknown) => string
+    descriptions: { list: string; add: string; update: string; remove: string; test: string }
+}
+
+export const registerConnectionTools = <TConfig extends { name: string }, TBackend extends Backend>(
+    server: McpServer,
+    deps: ConnectionToolsDeps<TConfig, TBackend>
+): void => {
+    const { registry, manager } = deps
+
+    server.registerTool('list_connections', { description: deps.descriptions.list }, async () =>
+        ok(registry.list().map(deps.publicView))
+    )
+
+    server.registerTool(
+        'add_connection',
+        {
+            description: deps.descriptions.add,
+            // Pass the strict object schema (not .shape) so the SDK rejects unknown
+            // keys at the boundary instead of silently stripping them.
+            inputSchema: deps.fullSchema
+        },
+        async (config) => {
+            try {
+                return ok({ added: deps.publicView(registry.add(deps.fullSchema.parse(config))) })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'update_connection',
+        {
+            description: deps.descriptions.update,
+            inputSchema: { name: z.string(), patch: deps.patchSchema }
+        },
+        async ({ name, patch }) => {
+            try {
+                const updated = registry.update(name, patch)
+                await manager.invalidate(name)
+                return ok({ updated: deps.publicView(updated) })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'remove_connection',
+        {
+            description: deps.descriptions.remove,
+            inputSchema: { name: z.string() }
+        },
+        async ({ name }) => {
+            try {
+                registry.remove(name)
+                await manager.invalidate(name)
+                return ok({ removed: name })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'test_connection',
+        {
+            description: deps.descriptions.test,
+            inputSchema: { name: z.string() }
+        },
+        async ({ name }) => {
+            // Start the clock before get() so cold tunnel/pool establishment counts.
+            const started = performance.now()
+            for (let attempt = 0; attempt < 2; attempt += 1) {
+                let managed: ManagedConnection<TConfig, TBackend>
+                try {
+                    managed = await manager.get(name)
+                } catch (error) {
+                    return fail(errorMessage(error))
+                }
+                try {
+                    const info = await deps.ping(managed)
+                    return ok({
+                        ok: true,
+                        ...info,
+                        latencyMs: Math.round(performance.now() - started)
+                    })
+                } catch (error) {
+                    // A concurrent reconfigure can dispose the backend under us;
+                    // rebuild once before reporting failure.
+                    if (error instanceof BackendDisposedError && attempt === 0) {
+                        continue
+                    }
+                    return fail(deps.formatError(managed.config, error))
+                }
+            }
+            return fail('connection is being reconfigured: retry the call')
+        }
+    )
+}

packages/core/src/mcp/managed.ts

@@ -0,0 +1,45 @@
+import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import type { Backend } from '../backend.js'
+import { BackendDisposedError } from '../backend.js'
+import type { ManagedConnection, Manager } from '../manager.js'
+import { errorMessage, fail, ok } from './respond.js'
+
+export type WithManagedDeps<TConfig> = {
+    isStaleError?: (error: unknown) => boolean
+    formatError: (config: TConfig, error: unknown) => string
+}
+
+// One transparent retry when the backend was disposed under the caller (a
+// concurrent invalidate/rotation ended its pools). manager.get on the retry
+// builds and returns the fresh entry.
+export const withManaged = async <TConfig extends { name: string }, TBackend extends Backend>(
+    manager: Manager<TConfig, TBackend>,
+    connection: string,
+    fn: (managed: ManagedConnection<TConfig, TBackend>) => Promise<unknown>,
+    deps: WithManagedDeps<TConfig>
+): Promise<CallToolResult> => {
+    const isStale = deps.isStaleError ?? ((e) => e instanceof BackendDisposedError)
+    for (let attempt = 0; attempt < 2; attempt += 1) {
+        let managed: ManagedConnection<TConfig, TBackend>
+        try {
+            managed = await manager.get(connection)
+        } catch (error) {
+            return fail(errorMessage(error))
+        }
+        try {
+            return ok(await fn(managed))
+        } catch (error) {
+            if (isStale(error)) {
+                if (attempt === 0) {
+                    continue
+                }
+                // Disposed again on the retry: the connection is being churned
+                // faster than we can use it. Surface a clean retryable message
+                // rather than the raw backend text.
+                break
+            }
+            return fail(deps.formatError(managed.config, error))
+        }
+    }
+    return fail('connection is being reconfigured: retry the call')
+}