From c7996a77ff7cc33a3c583c02608fa9a9f18d2555 Mon Sep 17 00:00:00 2001 From: anderson Date: Wed, 30 Sep 2026 11:21:42 -0300 Subject: [PATCH] feat(api): add PUT/DELETE /api/sources/:id with per-source reconnect Sources can be added, replaced, or removed at runtime without touching the other sources' connections. The TOML hot reload now applies the same per-source diff instead of disconnecting everything. Mutations require the HTTP bearer token to be configured and never touch sources defined by the startup configuration. --- .../sources-mutation.integration.test.ts | 176 ++++++++++++++++++ src/api/openapi.d.ts | 150 ++++++++++++++- src/api/openapi.yaml | 117 ++++++++++++ src/api/sources.ts | 175 ++++++++++++++++- src/config/toml-loader.ts | 10 + .../__tests__/manager-runtime-sources.test.ts | 73 ++++++++ src/connectors/manager.ts | 102 ++++++++++ src/server.ts | 12 +- src/tools/registry.ts | 56 +++++- src/utils/__tests__/config-watcher.test.ts | 162 +++++++++++----- src/utils/config-watcher.ts | 125 +++++++++---- src/utils/source-mutation.ts | 28 +++ 12 files changed, 1086 insertions(+), 100 deletions(-) create mode 100644 src/api/__tests__/sources-mutation.integration.test.ts create mode 100644 src/connectors/__tests__/manager-runtime-sources.test.ts create mode 100644 src/utils/source-mutation.ts diff --git a/src/api/__tests__/sources-mutation.integration.test.ts b/src/api/__tests__/sources-mutation.integration.test.ts new file mode 100644 index 00000000..dbbeddfb --- /dev/null +++ b/src/api/__tests__/sources-mutation.integration.test.ts @@ -0,0 +1,176 @@ +import { describe, it, expect, beforeAll, afterAll } from "vitest"; +import express, { Application } from "express"; +import { Server } from "http"; +import { setupManagerWithFixture, FIXTURES, loadFixtureConfig } from "../../__fixtures__/helpers.js"; +import type { ConnectorManager } from "../../connectors/manager.js"; +import { listSources, getSource, createSourceMutationHandlers } from "../sources.js"; +import { initializeToolRegistry, getToolRegistry } from "../../tools/registry.js"; +import { setFileSourceIds } from "../../utils/source-mutation.js"; +import type { components } from "../openapi.js"; + +// Import SQLite connector to ensure it's registered +import "../../connectors/sqlite/index.js"; + +type DataSource = components["schemas"]["DataSource"]; +type ErrorResponse = components["schemas"]["Error"]; + +const SECRET_DSN = "sqlite:///:memory:"; + +function buildApp(mutable: boolean): Application { + const app = express(); + app.use(express.json()); + const { putSource, deleteSource } = createSourceMutationHandlers({ mutable }); + app.get("/api/sources", listSources); + app.get("/api/sources/:sourceId", getSource); + app.put("/api/sources/:sourceId", (req, res) => { void putSource(req, res); }); + app.delete("/api/sources/:sourceId", (req, res) => { void deleteSource(req, res); }); + return app; +} + +async function listen(app: Application): Promise<{ server: Server; url: string }> { + const server = await new Promise((resolve) => { + const s = app.listen(0, () => resolve(s)); + }); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("no port"); + return { server, url: `http://localhost:${address.port}` }; +} + +describe("Sources API mutations", () => { + let manager: ConnectorManager; + let mutableServer: Server; + let immutableServer: Server; + let BASE_URL: string; + let LOCKED_URL: string; + + beforeAll(async () => { + manager = await setupManagerWithFixture(FIXTURES.MULTI_SQLITE); + const { sources, tools } = loadFixtureConfig(FIXTURES.MULTI_SQLITE); + initializeToolRegistry({ sources, tools: tools || [] }); + setFileSourceIds(sources.map((s) => s.id)); + + ({ server: mutableServer, url: BASE_URL } = await listen(buildApp(true))); + ({ server: immutableServer, url: LOCKED_URL } = await listen(buildApp(false))); + }, 30000); + + afterAll(async () => { + await new Promise((resolve) => mutableServer.close(() => resolve())); + await new Promise((resolve) => immutableServer.close(() => resolve())); + await manager.disconnect(); + }); + + it("adds a source and its default tools, without echoing the DSN", async () => { + const response = await fetch(`${BASE_URL}/api/sources/runtime_db`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ dsn: SECRET_DSN, description: "added at runtime", lazy: false, max_rows: 50 }), + }); + expect(response.status).toBe(200); + const raw = await response.text(); + expect(raw).not.toContain("sqlite:///"); + expect(JSON.parse(raw)).not.toHaveProperty("dsn"); + + const source = JSON.parse(raw) as DataSource; + expect(source.id).toBe("runtime_db"); + expect(source.type).toBe("sqlite"); + expect(source.description).toBe("added at runtime"); + const executeSql = source.tools.find((t) => t.name.startsWith("execute_sql")); + expect(executeSql?.readonly).toBe(true); + expect(executeSql?.max_rows).toBe(50); + + const registered = getToolRegistry().getEnabledToolConfigs("runtime_db").map((t) => t.name); + expect(registered).toEqual(["execute_sql", "search_objects"]); + }); + + it("lists the runtime source next to the file sources", async () => { + const response = await fetch(`${BASE_URL}/api/sources`); + const sources = (await response.json()) as DataSource[]; + expect(sources.map((s) => s.id)).toContain("runtime_db"); + expect(JSON.stringify(sources)).not.toContain("sqlite:///"); + }); + + it("replaces a runtime source in place", async () => { + const response = await fetch(`${BASE_URL}/api/sources/runtime_db`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ dsn: SECRET_DSN, description: "replaced", readonly: false }), + }); + expect(response.status).toBe(200); + const source = (await response.json()) as DataSource; + expect(source.description).toBe("replaced"); + expect(source.tools.find((t) => t.name.startsWith("execute_sql"))?.readonly).toBe(false); + + const all = (await (await fetch(`${BASE_URL}/api/sources`)).json()) as DataSource[]; + expect(all.filter((s) => s.id === "runtime_db")).toHaveLength(1); + }); + + it("keeps the file sources queryable while runtime sources change", async () => { + const result = await manager.getConnector("database_a").executeSQL("SELECT 1 AS one", {}); + expect(Number(result.resultSets[0].rows[0].one)).toBe(1); + }); + + it("removes the runtime source", async () => { + const response = await fetch(`${BASE_URL}/api/sources/runtime_db`, { method: "DELETE" }); + expect(response.status).toBe(204); + + const gone = await fetch(`${BASE_URL}/api/sources/runtime_db`); + expect(gone.status).toBe(404); + expect(getToolRegistry().getEnabledToolConfigs("runtime_db")).toEqual([]); + }); + + it("returns 404 when deleting an unknown source", async () => { + const response = await fetch(`${BASE_URL}/api/sources/never_existed`, { method: "DELETE" }); + expect(response.status).toBe(404); + }); + + it("refuses to change a source defined by the startup configuration", async () => { + const put = await fetch(`${BASE_URL}/api/sources/database_a`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ dsn: SECRET_DSN }), + }); + expect(put.status).toBe(409); + const del = await fetch(`${BASE_URL}/api/sources/database_a`, { method: "DELETE" }); + expect(del.status).toBe(409); + expect(manager.hasSource("database_a")).toBe(true); + }); + + it("rejects invalid bodies and ids", async () => { + const cases: Array<[string, unknown, string]> = [ + ["no_dsn", { description: "x" }, "'dsn' is required"], + ["array_body", [{ dsn: SECRET_DSN }], "JSON object"], + ["bad_scheme", { dsn: "ftp://nowhere/db" }, "not supported"], + ["bad_lazy", { dsn: SECRET_DSN, lazy: "yes" }, "'lazy' must be a boolean"], + ["bad_rows", { dsn: SECRET_DSN, max_rows: 0 }, "'max_rows' must be a positive integer"], + ]; + for (const [id, body, message] of cases) { + const response = await fetch(`${BASE_URL}/api/sources/${id}`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }); + expect(response.status, id).toBe(400); + const error = (await response.json()) as ErrorResponse; + expect(error.error, id).toContain(message); + } + + const badId = await fetch(`${BASE_URL}/api/sources/${encodeURIComponent("has space")}`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ dsn: SECRET_DSN }), + }); + expect(badId.status).toBe(400); + }); + + it("refuses mutations when the server has no bearer token", async () => { + const put = await fetch(`${LOCKED_URL}/api/sources/runtime_db`, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ dsn: SECRET_DSN }), + }); + expect(put.status).toBe(403); + const del = await fetch(`${LOCKED_URL}/api/sources/database_a`, { method: "DELETE" }); + expect(del.status).toBe(403); + expect(manager.hasSource("runtime_db")).toBe(false); + }); +}); diff --git a/src/api/openapi.d.ts b/src/api/openapi.d.ts index da3ca37d..adf69c03 100644 --- a/src/api/openapi.d.ts +++ b/src/api/openapi.d.ts @@ -36,9 +36,20 @@ export interface paths { * @description Returns details of a single data source by ID */ get: operations["getSource"]; - put?: never; + /** + * Add or replace a data source at runtime + * @description Creates the source, or replaces the source with the same id, and enables its + * built-in tools. Only this source is (re)connected; other sources are untouched. + * Requires bearer authentication to be configured on the server; sources defined + * by the startup configuration cannot be changed here. + */ + put: operations["putSource"]; post?: never; - delete?: never; + /** + * Remove a runtime data source + * @description Disconnects and forgets the source and its tools. Other sources are untouched. + */ + delete: operations["deleteSource"]; options?: never; head?: never; patch?: never; @@ -48,6 +59,31 @@ export interface paths { export type webhooks = Record; export interface components { schemas: { + SourceUpsert: { + /** + * @description Connection string; never returned by any response + * @example postgres://reader:secret@db.example.com:5432/app?sslmode=require + */ + dsn: string; + /** @description Human-readable description surfaced in tool descriptions */ + description?: string; + /** + * @description Connect on first use instead of during this request + * @default true + */ + lazy: boolean; + /** + * @description Restrict execute_sql to read-only statements + * @default true + */ + readonly: boolean; + /** @description Row cap for execute_sql SELECT results */ + max_rows?: number; + /** @description Per-query timeout in seconds */ + query_timeout?: number; + /** @description Connection timeout in seconds */ + connection_timeout?: number; + }; DataSource: { /** * @description Unique identifier for the data source @@ -257,4 +293,114 @@ export interface operations { }; }; }; + putSource: { + parameters: { + query?: never; + header?: never; + path: { + /** @description Source id; also the tool-name suffix */ + sourceId: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["SourceUpsert"]; + }; + }; + responses: { + /** @description The source as stored (credentials excluded) */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DataSource"]; + }; + }; + /** @description Invalid body or unsupported DSN */ + 400: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + /** @description Runtime source changes are disabled (no bearer token configured) */ + 403: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + /** @description The id belongs to a source from the startup configuration */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + /** @description The database could not be connected (eager sources only) */ + 502: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + }; + }; + deleteSource: { + parameters: { + query?: never; + header?: never; + path: { + sourceId: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Removed */ + 204: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Runtime source changes are disabled (no bearer token configured) */ + 403: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + /** @description Unknown source */ + 404: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + /** @description The id belongs to a source from the startup configuration */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Error"]; + }; + }; + }; + }; } diff --git a/src/api/openapi.yaml b/src/api/openapi.yaml index 293ed6fe..503c9fb4 100644 --- a/src/api/openapi.yaml +++ b/src/api/openapi.yaml @@ -82,9 +82,126 @@ paths: example: error: Source not found source_id: unknown-id + put: + summary: Add or replace a data source at runtime + description: | + Creates the source, or replaces the source with the same id, and enables its + built-in tools. Only this source is (re)connected; other sources are untouched. + Requires bearer authentication to be configured on the server; sources defined + by the startup configuration cannot be changed here. + operationId: putSource + parameters: + - name: sourceId + in: path + required: true + description: Source id; also the tool-name suffix + schema: + type: string + pattern: '^[A-Za-z0-9_-]{1,64}$' + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/SourceUpsert' + responses: + '200': + description: The source as stored (credentials excluded) + content: + application/json: + schema: + $ref: '#/components/schemas/DataSource' + '400': + description: Invalid body or unsupported DSN + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '403': + description: Runtime source changes are disabled (no bearer token configured) + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '409': + description: The id belongs to a source from the startup configuration + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '502': + description: The database could not be connected (eager sources only) + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + delete: + summary: Remove a runtime data source + description: Disconnects and forgets the source and its tools. Other sources are untouched. + operationId: deleteSource + parameters: + - name: sourceId + in: path + required: true + schema: + type: string + responses: + '204': + description: Removed + '403': + description: Runtime source changes are disabled (no bearer token configured) + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '404': + description: Unknown source + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '409': + description: The id belongs to a source from the startup configuration + content: + application/json: + schema: + $ref: '#/components/schemas/Error' components: schemas: + SourceUpsert: + type: object + required: + - dsn + properties: + dsn: + type: string + description: Connection string; never returned by any response + example: postgres://reader:secret@db.example.com:5432/app?sslmode=require + description: + type: string + description: Human-readable description surfaced in tool descriptions + lazy: + type: boolean + default: true + description: Connect on first use instead of during this request + readonly: + type: boolean + default: true + description: Restrict execute_sql to read-only statements + max_rows: + type: integer + minimum: 1 + description: Row cap for execute_sql SELECT results + query_timeout: + type: integer + minimum: 1 + description: Per-query timeout in seconds + connection_timeout: + type: integer + minimum: 1 + description: Connection timeout in seconds + DataSource: type: object required: diff --git a/src/api/sources.ts b/src/api/sources.ts index f3a9db1f..2d6ec4e6 100644 --- a/src/api/sources.ts +++ b/src/api/sources.ts @@ -1,13 +1,27 @@ import { Request, Response } from "express"; import { ConnectorManager } from "../connectors/manager.js"; +import { ConnectorRegistry } from "../connectors/interface.js"; +import { validateRuntimeSourceConfig } from "../config/toml-loader.js"; +import { redactDSN } from "../config/env.js"; +import { getToolRegistry } from "../tools/registry.js"; import { getDatabaseTypeFromDSN } from "../utils/dsn-obfuscate.js"; import { getToolsForSource } from "../utils/tool-metadata.js"; -import type { SourceConfig } from "../types/config.js"; +import { isFileSource, withSourceLock } from "../utils/source-mutation.js"; +import type { ExecuteSqlToolConfig, SearchObjectsToolConfig, SourceConfig, ToolConfig } from "../types/config.js"; import type { components } from "./openapi.js"; type DataSource = components["schemas"]["DataSource"]; type SSHTunnel = components["schemas"]["SSHTunnel"]; type ErrorResponse = components["schemas"]["Error"]; +type SourceUpsert = components["schemas"]["SourceUpsert"]; + +// Source ids double as tool-name suffixes, so they stay in the safe MCP tool-name alphabet. +const SOURCE_ID_PATTERN = /^[A-Za-z0-9_-]{1,64}$/; + +export interface SourceMutationOptions { + /** Runtime source changes are only allowed when the HTTP transport requires a bearer token. */ + mutable: boolean; +} /** * Transform a SourceConfig into an API DataSource response @@ -29,6 +43,7 @@ function transformSourceConfig(source: SourceConfig): DataSource { const dataSource: DataSource = { id: source.id, type: source.type, + tools: [], }; // Add description if present @@ -126,3 +141,161 @@ export function getSource(req: Request, res: Response): void { res.status(500).json(errorResponse); } } + +/** + * PUT /api/sources/:sourceId and DELETE /api/sources/:sourceId + * Add, replace, or remove one source at runtime without touching the others. + * Sources defined by the startup configuration are immutable here (409). + */ +export function createSourceMutationHandlers(options: SourceMutationOptions): { + putSource: (req: Request, res: Response) => Promise; + deleteSource: (req: Request, res: Response) => Promise; +} { + const guard = (sourceId: string, res: Response): boolean => { + if (!options.mutable) { + res.status(403).json({ + error: "Runtime source changes require a bearer token (--auth-token / DBHUB_AUTH_TOKEN)", + source_id: sourceId, + } satisfies ErrorResponse); + return false; + } + if (!SOURCE_ID_PATTERN.test(sourceId)) { + res.status(400).json({ + error: "Source id must match [A-Za-z0-9_-]{1,64}", + source_id: sourceId, + } satisfies ErrorResponse); + return false; + } + if (isFileSource(sourceId)) { + res.status(409).json({ + error: "Source is defined by the startup configuration and cannot be changed over the API", + source_id: sourceId, + } satisfies ErrorResponse); + return false; + } + return true; + }; + + const putSource = async (req: Request, res: Response): Promise => { + const sourceId = req.params.sourceId; + if (!guard(sourceId, res)) { + return; + } + + let source: SourceConfig; + let tools: ToolConfig[]; + try { + ({ source, tools } = buildRuntimeSource(sourceId, req.body)); + } catch (error) { + res.status(400).json({ + error: error instanceof Error ? error.message : "Invalid source", + source_id: sourceId, + } satisfies ErrorResponse); + return; + } + + try { + await withSourceLock(async () => { + await ConnectorManager.addSource(source); + getToolRegistry().setSourceTools(sourceId, tools); + }); + } catch (error) { + // A failed add leaves no source behind; a failed replace already removed the old one. + getToolRegistry().removeSource(sourceId); + const message = error instanceof Error ? error.message : "Failed to add source"; + res.status(502).json({ + error: message.split(source.dsn!).join(redactDSN(source.dsn!)), + source_id: sourceId, + } satisfies ErrorResponse); + return; + } + + console.error(`Source '${sourceId}' ${source.lazy ? "registered" : "connected"} via API`); + res.json(transformSourceConfig(ConnectorManager.getSourceConfig(sourceId)!)); + }; + + const deleteSource = async (req: Request, res: Response): Promise => { + const sourceId = req.params.sourceId; + if (!guard(sourceId, res)) { + return; + } + + const removed = await withSourceLock(async () => { + const found = await ConnectorManager.removeSource(sourceId); + if (found) { + getToolRegistry().removeSource(sourceId); + } + return found; + }); + if (!removed) { + res.status(404).json({ error: "Source not found", source_id: sourceId } satisfies ErrorResponse); + return; + } + + console.error(`Source '${sourceId}' removed via API`); + res.status(204).end(); + }; + + return { putSource, deleteSource }; +} + +/** + * Turn a PUT body into a validated SourceConfig plus its two built-in tools. + * Throws with a message safe to return to the caller (never echoes the DSN). + */ +function buildRuntimeSource(sourceId: string, body: unknown): { source: SourceConfig; tools: ToolConfig[] } { + if (!body || typeof body !== "object" || Array.isArray(body)) { + throw new Error("Request body must be a JSON object"); + } + const input = body as Partial; + + if (typeof input.dsn !== "string" || input.dsn.trim() === "") { + throw new Error("'dsn' is required"); + } + const dsn = input.dsn.trim(); + const type = getDatabaseTypeFromDSN(dsn); + if (!type || !ConnectorRegistry.getConnectorForDSN(dsn)) { + throw new Error("'dsn' scheme is not supported"); + } + + const optionalString = (key: "description"): string | undefined => { + const value = input[key]; + if (value === undefined || value === null) return undefined; + if (typeof value !== "string") throw new Error(`'${key}' must be a string`); + return value; + }; + const optionalBoolean = (key: "lazy" | "readonly", fallback: boolean): boolean => { + const value = input[key]; + if (value === undefined || value === null) return fallback; + if (typeof value !== "boolean") throw new Error(`'${key}' must be a boolean`); + return value; + }; + const optionalPositiveInt = (key: "max_rows" | "query_timeout" | "connection_timeout"): number | undefined => { + const value = input[key]; + if (value === undefined || value === null) return undefined; + if (typeof value !== "number" || !Number.isInteger(value) || value < 1) { + throw new Error(`'${key}' must be a positive integer`); + } + return value; + }; + + const source: SourceConfig = { + id: sourceId, + dsn, + type, + description: optionalString("description"), + lazy: optionalBoolean("lazy", true), + query_timeout: optionalPositiveInt("query_timeout"), + connection_timeout: optionalPositiveInt("connection_timeout"), + }; + validateRuntimeSourceConfig(source); + + const executeSql: ExecuteSqlToolConfig = { + name: "execute_sql", + source: sourceId, + readonly: optionalBoolean("readonly", true), + max_rows: optionalPositiveInt("max_rows"), + }; + const searchObjects: SearchObjectsToolConfig = { name: "search_objects", source: sourceId }; + return { source, tools: [executeSql, searchObjects] }; +} diff --git a/src/config/toml-loader.ts b/src/config/toml-loader.ts index 8d285c49..b42ea9fd 100644 --- a/src/config/toml-loader.ts +++ b/src/config/toml-loader.ts @@ -923,6 +923,16 @@ function processSourceConfigs( * Unresolved variables are left as-is (no error thrown). */ const ENV_VAR_PATTERN = /\$\{([^}]+)\}/g; +/** + * Validate a single source supplied at runtime (sources API). Same rules as a + * TOML [[sources]] entry; the "file" in messages is the API origin. + */ +export function validateRuntimeSourceConfig(source: SourceConfig): void { + const origin = "sources API"; + validateSourceConfig(source, origin); + validateDSNFieldConflicts(source, origin); +} + export function interpolateEnvVars(value: unknown): unknown { if (typeof value === "string") { return value.replace(ENV_VAR_PATTERN, (match, varName) => { diff --git a/src/connectors/__tests__/manager-runtime-sources.test.ts b/src/connectors/__tests__/manager-runtime-sources.test.ts new file mode 100644 index 00000000..ef2aa320 --- /dev/null +++ b/src/connectors/__tests__/manager-runtime-sources.test.ts @@ -0,0 +1,73 @@ +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import { ConnectorManager } from "../manager.js"; +import "../sqlite/index.js"; + +const memory = (id: string, lazy = false) => ({ id, type: "sqlite" as const, dsn: "sqlite:///:memory:", lazy }); + +describe("ConnectorManager runtime sources", () => { + let manager: ConnectorManager; + + beforeEach(async () => { + manager = new ConnectorManager(); + await manager.connectWithSources([memory("base")]); + }); + + afterEach(async () => { + await manager.disconnect(); + }); + + it("adds an eager source and leaves the existing one connected", async () => { + await manager.addSource(memory("extra")); + + expect(manager.getSourceIds()).toEqual(["base", "extra"]); + expect(manager.hasSource("extra")).toBe(true); + const result = await manager.getConnector("extra").executeSQL("SELECT 1 AS one", {}); + expect(Number(result.resultSets[0].rows[0].one)).toBe(1); + const base = await manager.getConnector("base").executeSQL("SELECT 2 AS two", {}); + expect(Number(base.resultSets[0].rows[0].two)).toBe(2); + }); + + it("registers a lazy source without connecting until first use", async () => { + await manager.addSource(memory("later", true)); + + expect(manager.hasSource("later")).toBe(true); + expect(() => manager.getConnector("later")).toThrow(/not found/); + await manager.ensureConnected("later"); + const result = await manager.getConnector("later").executeSQL("SELECT 3 AS three", {}); + expect(Number(result.resultSets[0].rows[0].three)).toBe(3); + }); + + it("removes only the named source", async () => { + await manager.addSource(memory("extra")); + + expect(await manager.removeSource("extra")).toBe(true); + expect(manager.getSourceIds()).toEqual(["base"]); + expect(manager.hasSource("extra")).toBe(false); + expect(manager.getSourceConfig("extra")).toBeNull(); + const base = await manager.getConnector("base").executeSQL("SELECT 4 AS four", {}); + expect(Number(base.resultSets[0].rows[0].four)).toBe(4); + }); + + it("reports false for an unknown id", async () => { + expect(await manager.removeSource("nope")).toBe(false); + expect(manager.getSourceIds()).toEqual(["base"]); + }); + + it("replaces a source with the same id in place", async () => { + await manager.addSource({ ...memory("extra"), description: "first" }); + await manager.addSource({ ...memory("extra"), description: "second" }); + + expect(manager.getSourceIds()).toEqual(["base", "extra"]); + expect(manager.getSourceConfig("extra")?.description).toBe("second"); + const result = await manager.getConnector("extra").executeSQL("SELECT 5 AS five", {}); + expect(Number(result.resultSets[0].rows[0].five)).toBe(5); + }); + + it("removes a lazy source that never connected", async () => { + await manager.addSource(memory("later", true)); + + expect(await manager.removeSource("later")).toBe(true); + expect(manager.hasSource("later")).toBe(false); + await expect(manager.ensureConnected("later")).rejects.toThrow(/not found/); + }); +}); diff --git a/src/connectors/manager.ts b/src/connectors/manager.ts index 1b4c309a..21b2640b 100644 --- a/src/connectors/manager.ts +++ b/src/connectors/manager.ts @@ -360,6 +360,108 @@ export class ConnectorManager { this.isDisconnecting = false; } + /** + * Add a source at runtime, or replace the source with the same id. + * Only this source is touched: other sources keep their pools and tunnels. + * A lazy source is registered without connecting; an eager one connects now + * and the error propagates (nothing is registered on failure). + */ + async addSource(source: SourceConfig): Promise { + if (this.hasSource(source.id)) { + await this.removeSource(source.id); + } + if (source.lazy) { + this.registerLazySource(source); + return; + } + await this.connectSource(source); + } + + /** + * Remove one source at runtime: stop its IAM refresh, wait for any in-flight + * lazy connect, close its connector and tunnel, and drop it from every map. + * Returns false when the id is unknown. Other sources are untouched. + */ + async removeSource(sourceId: string): Promise { + if (!this.hasSource(sourceId)) { + return false; + } + + const timer = this.iamRefreshTimers.get(sourceId); + if (timer) { + clearTimeout(timer); + this.iamRefreshTimers.delete(sourceId); + } + + // A lazy connect in progress would otherwise re-register the connector after removal. + const pending = this.pendingConnections.get(sourceId); + if (pending) { + try { + await pending; + } catch { + // The connect failed on its own; there is nothing left to close. + } + } + + const connector = this.connectors.get(sourceId); + if (connector) { + try { + await connector.disconnect(); + console.error(`Disconnected from source '${sourceId}'`); + } catch (error) { + console.error(`Error disconnecting from source '${sourceId}':`, error); + } + } + + const tunnel = this.sshTunnels.get(sourceId); + if (tunnel) { + try { + await tunnel.close(); + } catch (error) { + console.error(`Error closing SSH tunnel for source '${sourceId}':`, error); + } + } + + this.connectors.delete(sourceId); + this.sshTunnels.delete(sourceId); + this.sourceConfigs.delete(sourceId); + this.lazySources.delete(sourceId); + this.pendingConnections.delete(sourceId); + this.sourceIds = this.sourceIds.filter((id) => id !== sourceId); + return true; + } + + /** + * Whether a source id is registered (connected or pending lazy connection) + */ + hasSource(sourceId: string): boolean { + return this.sourceConfigs.has(sourceId); + } + + /** Runtime add/replace of a source (static access for API handlers) */ + static async addSource(source: SourceConfig): Promise { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.addSource(source); + } + + /** Runtime removal of a source (static access for API handlers) */ + static async removeSource(sourceId: string): Promise { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.removeSource(sourceId); + } + + /** Whether a source id is registered (static access for API handlers) */ + static hasSource(sourceId: string): boolean { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.hasSource(sourceId); + } + /** * Get a connector by source ID * If sourceId is not provided, returns the default (first) connector diff --git a/src/server.ts b/src/server.ts index c9585b58..864c6756 100644 --- a/src/server.ts +++ b/src/server.ts @@ -11,11 +11,12 @@ import { ConnectorManager } from "./connectors/manager.js"; import { ConnectorRegistry } from "./connectors/interface.js"; import { resolveTransport, resolvePort, resolveHost, resolveAllowedHosts, resolveAuthTokens, resolveSourceConfigs, isDemoMode } from "./config/env.js"; import { registerTools } from "./tools/index.js"; -import { listSources, getSource } from "./api/sources.js"; +import { listSources, getSource, createSourceMutationHandlers } from "./api/sources.js"; import { listRequests } from "./api/requests.js"; import { generateStartupTable, buildSourceDisplayInfo } from "./utils/startup-table.js"; import { getToolsForSource } from "./utils/tool-metadata.js"; import { startConfigWatcher } from "./utils/config-watcher.js"; +import { setFileSourceIds } from "./utils/source-mutation.js"; import { validateOrigin, buildAllowedHosts, getSelfHosts, ALLOW_ANY_HOST } from "./utils/cross-origin.js"; import { validateAuthToken } from "./utils/auth-token.js"; @@ -105,6 +106,8 @@ See documentation for more details on configuring database connections. // Connect to database(s) - works uniformly for all modes (demo, single DSN, multi-source TOML) await connectorManager.connectWithSources(sources); + // Startup sources are owned by the configuration; the sources API may not alter them. + setFileSourceIds(sources.map((s) => s.id)); // Initialize tool registry (manages both built-in and custom tools) // This must happen AFTER ConnectorManager is initialized so source validation works @@ -230,7 +233,7 @@ See documentation for more details on configuring database connections. // named and cannot be statically allow-listed — browser clients skip // mirroring them for exactly this reason.) res.header('Access-Control-Allow-Origin', origin || 'http://localhost'); - res.header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS'); + res.header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS'); res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization, Mcp-Session-Id, MCP-Protocol-Version, Mcp-Method, Mcp-Name'); res.header('Access-Control-Allow-Credentials', 'true'); @@ -264,9 +267,12 @@ See documentation for more details on configuring database connections. const frontendPath = path.join(__dirname, "public"); app.use(express.static(frontendPath)); - // Data sources API endpoints + // Data sources API endpoints. Mutations are enabled only behind bearer auth. + const { putSource, deleteSource } = createSourceMutationHandlers({ mutable: authTokens.length > 0 }); app.get("/api/sources", listSources); app.get("/api/sources/:sourceId", getSource); + app.put("/api/sources/:sourceId", (req, res) => { void putSource(req, res); }); + app.delete("/api/sources/:sourceId", (req, res) => { void deleteSource(req, res); }); app.get("/api/requests", listRequests); // Main MCP endpoint. createMcpHandler serves both protocol eras from diff --git a/src/tools/registry.ts b/src/tools/registry.ts index d519c2c8..8ffc68d5 100644 --- a/src/tools/registry.ts +++ b/src/tools/registry.ts @@ -183,21 +183,59 @@ export class ToolRegistry { // Backward compatibility: sources without tools get default built-ins for (const source of config.sources) { if (!registry.has(source.id)) { - const defaultTools: ToolConfig[] = BUILTIN_TOOLS.map((name) => { - // Create properly typed tool configs based on the tool name - if (name === 'execute_sql') { - return { name: 'execute_sql', source: source.id } satisfies ExecuteSqlToolConfig; - } else { - return { name: 'search_objects', source: source.id } satisfies SearchObjectsToolConfig; - } - }); - registry.set(source.id, defaultTools); + registry.set(source.id, this.defaultToolsFor(source.id)); } } return registry; } + /** + * Set the enabled tools for one source at runtime (added or replaced source). + * An empty list enables the default built-ins, matching TOML behavior. + * Custom tools are validated against the sources currently registered. + */ + setSourceTools(sourceId: string, tools: ToolConfig[]): void { + const availableSources = [...new Set([...this.toolsBySource.keys(), sourceId])]; + const otherCustomNames = new Set( + [...this.toolsBySource.entries()] + .filter(([id]) => id !== sourceId) + .flatMap(([, list]) => list) + .filter((t) => !this.isBuiltinTool(t.name)) + .map((t) => t.name) + ); + for (const tool of tools) { + if (tool.source !== sourceId) { + throw new Error(`Tool '${tool.name}' targets source '${tool.source}', expected '${sourceId}'`); + } + if (!this.isBuiltinTool(tool.name)) { + this.validateCustomTool(tool, availableSources); + if (otherCustomNames.has(tool.name)) { + throw new Error(`Duplicate tool name '${tool.name}'. Tool names must be unique.`); + } + } + } + this.toolsBySource.set(sourceId, tools.length > 0 ? [...tools] : this.defaultToolsFor(sourceId)); + } + + /** + * Forget one source's tools (removed source). No-op for unknown ids. + */ + removeSource(sourceId: string): void { + this.toolsBySource.delete(sourceId); + } + + /** + * Default built-in tool configs for a source that declares none + */ + private defaultToolsFor(sourceId: string): ToolConfig[] { + return BUILTIN_TOOLS.map((name) => + name === "execute_sql" + ? ({ name: "execute_sql", source: sourceId } satisfies ExecuteSqlToolConfig) + : ({ name: "search_objects", source: sourceId } satisfies SearchObjectsToolConfig) + ); + } + /** * Get all enabled tool configs for a specific source */ diff --git a/src/utils/__tests__/config-watcher.test.ts b/src/utils/__tests__/config-watcher.test.ts index 2043ed72..ec6934c6 100644 --- a/src/utils/__tests__/config-watcher.test.ts +++ b/src/utils/__tests__/config-watcher.test.ts @@ -10,16 +10,16 @@ vi.mock("../../config/toml-loader.js", () => ({ loadTomlConfig: vi.fn(), })); vi.mock("../../tools/registry.js", () => ({ - initializeToolRegistry: vi.fn(), + getToolRegistry: vi.fn(), })); import { resolveTomlConfigPath, loadTomlConfig } from "../../config/toml-loader.js"; -import { initializeToolRegistry } from "../../tools/registry.js"; +import { getToolRegistry } from "../../tools/registry.js"; function createMockManager(overrides: Partial> = {}) { return { - disconnect: vi.fn().mockResolvedValue(undefined), - connectWithSources: vi.fn().mockResolvedValue(undefined), + addSource: vi.fn().mockResolvedValue(undefined), + removeSource: vi.fn().mockResolvedValue(true), getAllSourceConfigs: vi.fn().mockReturnValue([]), ...overrides, } as unknown as ConnectorManager; @@ -29,11 +29,16 @@ function createOptions(connectorManager: ConnectorManager, initialTools?: any[]) return { connectorManager, initialTools }; } +const OLD_DB = { id: "old_db", type: "sqlite" as const, dsn: "sqlite:///:memory:" }; +const OLD_TOOLS = [{ name: "execute_sql" as const, source: "old_db" }]; + describe("startConfigWatcher", () => { let mockWatcher: { on: ReturnType; close: ReturnType; unref: ReturnType }; let watchCallback: (eventType: string) => void; + let registry: { setSourceTools: ReturnType; removeSource: ReturnType }; beforeEach(() => { + vi.clearAllMocks(); vi.useFakeTimers(); mockWatcher = { on: vi.fn().mockReturnThis(), @@ -44,6 +49,8 @@ describe("startConfigWatcher", () => { watchCallback = cb; return mockWatcher as any; }); + registry = { setSourceTools: vi.fn(), removeSource: vi.fn() }; + vi.mocked(getToolRegistry).mockReturnValue(registry as any); }); afterEach(() => { @@ -68,32 +75,74 @@ describe("startConfigWatcher", () => { expect(mockWatcher.unref).toHaveBeenCalled(); }); - it("should reload config on file change after debounce", async () => { + it("should add new sources on file change after debounce", async () => { vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); - const newConfig = { - sources: [{ id: "new_db", type: "postgres" as const, dsn: "postgres://localhost/new" }], - tools: [], - source: "dbhub.toml", - }; - vi.mocked(loadTomlConfig).mockReturnValue(newConfig); + const newSource = { id: "new_db", type: "postgres" as const, dsn: "postgres://localhost/new" }; + const newTools = [{ name: "execute_sql" as const, source: "new_db", readonly: true }]; + vi.mocked(loadTomlConfig).mockReturnValue({ sources: [newSource], tools: newTools, source: "dbhub.toml" }); const mockManager = createMockManager(); startConfigWatcher(createOptions(mockManager)); watchCallback("change"); // Before debounce, nothing should happen - expect(mockManager.disconnect).not.toHaveBeenCalled(); + expect(mockManager.addSource).not.toHaveBeenCalled(); // After debounce await vi.advanceTimersByTimeAsync(500); expect(loadTomlConfig).toHaveBeenCalled(); - expect(mockManager.disconnect).toHaveBeenCalled(); - expect(mockManager.connectWithSources).toHaveBeenCalledWith(newConfig.sources); - expect(initializeToolRegistry).toHaveBeenCalledWith({ - sources: newConfig.sources, - tools: newConfig.tools, + expect(mockManager.addSource).toHaveBeenCalledWith(newSource); + expect(registry.setSourceTools).toHaveBeenCalledWith("new_db", newTools); + expect(mockManager.removeSource).not.toHaveBeenCalled(); + }); + + it("should leave unchanged sources untouched", async () => { + vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); + const added = { id: "added_db", type: "sqlite" as const, dsn: "sqlite:///:memory:" }; + vi.mocked(loadTomlConfig).mockReturnValue({ + sources: [{ ...OLD_DB }, added], + tools: [...OLD_TOOLS], + source: "dbhub.toml", }); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); + + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); + watchCallback("change"); + await vi.advanceTimersByTimeAsync(500); + + expect(mockManager.addSource).toHaveBeenCalledTimes(1); + expect(mockManager.addSource).toHaveBeenCalledWith(added); + expect(mockManager.removeSource).not.toHaveBeenCalled(); + expect(registry.setSourceTools).not.toHaveBeenCalledWith("old_db", expect.anything()); + }); + + it("should remove sources that disappeared from the file", async () => { + vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); + vi.mocked(loadTomlConfig).mockReturnValue({ sources: [], tools: [], source: "dbhub.toml" }); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); + + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); + watchCallback("change"); + await vi.advanceTimersByTimeAsync(500); + + expect(mockManager.removeSource).toHaveBeenCalledWith("old_db"); + expect(registry.removeSource).toHaveBeenCalledWith("old_db"); + expect(mockManager.addSource).not.toHaveBeenCalled(); + }); + + it("should re-add a source whose tools changed", async () => { + vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); + const changedTools = [{ name: "execute_sql" as const, source: "old_db", readonly: true }]; + vi.mocked(loadTomlConfig).mockReturnValue({ sources: [{ ...OLD_DB }], tools: changedTools, source: "dbhub.toml" }); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); + + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); + watchCallback("change"); + await vi.advanceTimersByTimeAsync(500); + + expect(mockManager.addSource).toHaveBeenCalledWith(OLD_DB); + expect(registry.setSourceTools).toHaveBeenCalledWith("old_db", changedTools); }); it("should debounce rapid file changes", async () => { @@ -111,7 +160,8 @@ describe("startConfigWatcher", () => { watchCallback("change"); await vi.advanceTimersByTimeAsync(500); - expect(mockManager.disconnect).toHaveBeenCalledTimes(1); + expect(loadTomlConfig).toHaveBeenCalledTimes(1); + expect(mockManager.addSource).toHaveBeenCalledTimes(1); }); it("should keep existing connections when new config is invalid", async () => { @@ -119,74 +169,88 @@ describe("startConfigWatcher", () => { vi.mocked(loadTomlConfig).mockImplementation(() => { throw new Error("Invalid TOML"); }); - const mockManager = createMockManager(); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); - startConfigWatcher(createOptions(mockManager)); + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); watchCallback("change"); await vi.advanceTimersByTimeAsync(500); - expect(mockManager.disconnect).not.toHaveBeenCalled(); + expect(mockManager.removeSource).not.toHaveBeenCalled(); + expect(mockManager.addSource).not.toHaveBeenCalled(); }); it("should keep existing connections when loadTomlConfig returns null", async () => { vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); vi.mocked(loadTomlConfig).mockReturnValue(null); - const mockManager = createMockManager(); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); - startConfigWatcher(createOptions(mockManager)); + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); watchCallback("change"); await vi.advanceTimersByTimeAsync(500); - expect(mockManager.disconnect).not.toHaveBeenCalled(); + expect(mockManager.removeSource).not.toHaveBeenCalled(); + expect(mockManager.addSource).not.toHaveBeenCalled(); }); - it("should rollback with initial tools when connectWithSources fails", async () => { + it("should roll a changed source back to its previous definition when the new one fails", async () => { vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); - const newConfig = { - sources: [{ id: "bad_db", type: "postgres" as const, dsn: "postgres://localhost/bad" }], - tools: [{ name: "execute_sql" as const, source: "bad_db", readonly: true }], - source: "dbhub.toml", - }; - vi.mocked(loadTomlConfig).mockReturnValue(newConfig); + const badSource = { id: "old_db", type: "postgres" as const, dsn: "postgres://localhost/bad" }; + const badTools = [{ name: "execute_sql" as const, source: "old_db", readonly: true }]; + vi.mocked(loadTomlConfig).mockReturnValue({ sources: [badSource], tools: badTools, source: "dbhub.toml" }); - const oldSources = [{ id: "old_db", type: "sqlite" as const, dsn: "sqlite:///:memory:" }]; - const oldTools = [{ name: "execute_sql" as const, source: "old_db" }]; const mockManager = createMockManager({ - connectWithSources: vi.fn() + addSource: vi.fn() .mockRejectedValueOnce(new Error("Connection refused")) .mockResolvedValueOnce(undefined), - getAllSourceConfigs: vi.fn().mockReturnValue(oldSources), + getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]), }); - startConfigWatcher(createOptions(mockManager, oldTools)); + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); watchCallback("change"); await vi.advanceTimersByTimeAsync(500); - expect(mockManager.connectWithSources).toHaveBeenNthCalledWith(1, newConfig.sources); - expect(mockManager.connectWithSources).toHaveBeenLastCalledWith(oldSources); - expect(initializeToolRegistry).toHaveBeenLastCalledWith({ sources: oldSources, tools: oldTools }); + expect(mockManager.addSource).toHaveBeenNthCalledWith(1, badSource); + expect(mockManager.addSource).toHaveBeenLastCalledWith(OLD_DB); + expect(registry.setSourceTools).toHaveBeenLastCalledWith("old_db", OLD_TOOLS); }); - it("should disconnect partial state before rollback", async () => { + it("should drop a brand-new source that fails instead of leaving partial state", async () => { vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); vi.mocked(loadTomlConfig).mockReturnValue({ - sources: [{ id: "bad", type: "postgres" as const, dsn: "postgres://localhost/bad" }], - tools: [], + sources: [{ ...OLD_DB }, { id: "bad", type: "postgres" as const, dsn: "postgres://localhost/bad" }], + tools: [...OLD_TOOLS], source: "dbhub.toml", }); - const mockManager = createMockManager({ - connectWithSources: vi.fn() - .mockRejectedValueOnce(new Error("Partial failure")) - .mockResolvedValueOnce(undefined), + addSource: vi.fn().mockRejectedValueOnce(new Error("Partial failure")), + getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]), }); - startConfigWatcher(createOptions(mockManager)); + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); watchCallback("change"); await vi.advanceTimersByTimeAsync(500); - // disconnect called twice: once for initial teardown, once to clean up partial state before rollback - expect(mockManager.disconnect).toHaveBeenCalledTimes(2); + expect(mockManager.addSource).toHaveBeenCalledTimes(1); + expect(registry.removeSource).toHaveBeenCalledWith("bad"); + expect(mockManager.removeSource).not.toHaveBeenCalledWith("old_db"); + }); + + it("should treat the effective config as the baseline for the next reload", async () => { + vi.mocked(resolveTomlConfigPath).mockReturnValue("/path/to/dbhub.toml"); + const added = { id: "added_db", type: "sqlite" as const, dsn: "sqlite:///:memory:" }; + vi.mocked(loadTomlConfig).mockReturnValue({ sources: [{ ...OLD_DB }, added], tools: [...OLD_TOOLS], source: "dbhub.toml" }); + const mockManager = createMockManager({ getAllSourceConfigs: vi.fn().mockReturnValue([OLD_DB]) }); + + startConfigWatcher(createOptions(mockManager, OLD_TOOLS)); + watchCallback("change"); + await vi.advanceTimersByTimeAsync(500); + expect(mockManager.addSource).toHaveBeenCalledTimes(1); + + // Same file again: nothing to do + watchCallback("change"); + await vi.advanceTimersByTimeAsync(500); + expect(mockManager.addSource).toHaveBeenCalledTimes(1); + expect(mockManager.removeSource).not.toHaveBeenCalled(); }); it("should clean up watcher on cleanup call", () => { diff --git a/src/utils/config-watcher.ts b/src/utils/config-watcher.ts index 4f102391..fc5be31a 100644 --- a/src/utils/config-watcher.ts +++ b/src/utils/config-watcher.ts @@ -1,7 +1,8 @@ import fs from "fs"; import { loadTomlConfig, resolveTomlConfigPath } from "../config/toml-loader.js"; import { ConnectorManager } from "../connectors/manager.js"; -import { initializeToolRegistry } from "../tools/registry.js"; +import { getToolRegistry } from "../tools/registry.js"; +import { setFileSourceIds, withSourceLock } from "./source-mutation.js"; import type { SourceConfig, ToolConfig } from "../types/config.js"; const DEBOUNCE_MS = 500; @@ -15,6 +16,11 @@ interface ConfigWatcherOptions { * Watch the TOML configuration file for changes and reload sources automatically. * Only applicable when using TOML-based configuration. * + * A reload is a diff against the last known-good file: sources that are unchanged + * keep their connections and in-flight queries; only added, removed, or edited + * sources (or sources whose tools changed) are touched. Sources added through the + * sources API are not part of the file and are never affected by a reload. + * * NOTE: In STDIO transport mode, the MCP server's tool list is registered once at * startup. Hot reload updates the underlying database connections and tool registry, * but STDIO clients won't see added/removed tools until a full server restart. @@ -33,7 +39,7 @@ export function startConfigWatcher(options: ConfigWatcherOptions): (() => void) // Track last known-good config for rollback (sources + tools) let lastGoodSources: SourceConfig[] = connectorManager.getAllSourceConfigs(); - let lastGoodTools: ToolConfig[] | undefined = initialTools; + let lastGoodTools: ToolConfig[] = initialTools ?? []; const scheduleReload = () => { if (debounceTimer) { @@ -60,40 +66,18 @@ export function startConfigWatcher(options: ConfigWatcherOptions): (() => void) return; } - // Save current config for rollback - const oldSources = lastGoodSources; - const oldTools = lastGoodTools; - - // Disconnect all existing sources - await connectorManager.disconnect(); - - try { - // Reconnect with new sources - await connectorManager.connectWithSources(newConfig.sources); - - // Re-initialize tool registry with new config - initializeToolRegistry({ - sources: newConfig.sources, - tools: newConfig.tools, - }); - - // Update last known-good config - lastGoodSources = newConfig.sources; - lastGoodTools = newConfig.tools; - - console.error("Configuration reloaded successfully."); - } catch (connectError) { - console.error("Failed to connect with new config, rolling back:", connectError); - // Clean up any partial connections before rollback - try { await connectorManager.disconnect(); } catch { /* best effort */ } - try { - await connectorManager.connectWithSources(oldSources); - initializeToolRegistry({ sources: oldSources, tools: oldTools }); - console.error("Rolled back to previous configuration."); - } catch (rollbackError) { - console.error("Rollback also failed, server has no active connections:", rollbackError); - } - } + await withSourceLock(async () => { + const result = await applySourceDiff( + connectorManager, + { sources: lastGoodSources, tools: lastGoodTools }, + { sources: newConfig.sources, tools: newConfig.tools ?? [] } + ); + lastGoodSources = result.sources; + lastGoodTools = result.tools; + setFileSourceIds(lastGoodSources.map((s) => s.id)); + }); + + console.error("Configuration reloaded successfully."); } catch (error) { console.error("Config reload failed, keeping existing connections:", error); } finally { @@ -125,3 +109,72 @@ export function startConfigWatcher(options: ConfigWatcherOptions): (() => void) watcher.close(); }; } + +interface FileConfig { + sources: SourceConfig[]; + tools: ToolConfig[]; +} + +/** + * Apply `next` on top of `current` one source at a time. Returns the config that + * is actually in effect afterwards: a source whose new definition fails to connect + * is rolled back to its previous definition (or dropped, if it had none). + */ +async function applySourceDiff( + connectorManager: ConnectorManager, + current: FileConfig, + next: FileConfig +): Promise { + const registry = getToolRegistry(); + const currentById = new Map(current.sources.map((s) => [s.id, s])); + const nextById = new Map(next.sources.map((s) => [s.id, s])); + const toolsFor = (id: string, tools: ToolConfig[]) => tools.filter((t) => t.source === id); + const unchanged = (id: string) => + JSON.stringify(currentById.get(id)) === JSON.stringify(nextById.get(id)) && + JSON.stringify(toolsFor(id, current.tools)) === JSON.stringify(toolsFor(id, next.tools)); + + const effective = new Map(); + for (const source of current.sources) { + effective.set(source.id, { source, tools: toolsFor(source.id, current.tools) }); + } + + for (const id of currentById.keys()) { + if (nextById.has(id)) continue; + await connectorManager.removeSource(id); + registry.removeSource(id); + effective.delete(id); + console.error(`Source '${id}' removed`); + } + + for (const [id, source] of nextById) { + if (currentById.has(id) && unchanged(id)) continue; + const tools = toolsFor(id, next.tools); + try { + await connectorManager.addSource(source); + registry.setSourceTools(id, tools); + effective.set(id, { source, tools }); + console.error(`Source '${id}' ${currentById.has(id) ? "updated" : "added"}`); + } catch (error) { + console.error(`Source '${id}': failed to apply new config, rolling back:`, error); + const previous = effective.get(id); + if (!previous) { + registry.removeSource(id); + continue; + } + try { + await connectorManager.addSource(previous.source); + registry.setSourceTools(id, previous.tools); + console.error(`Source '${id}' rolled back to previous configuration.`); + } catch (rollbackError) { + console.error(`Source '${id}': rollback also failed, source is unavailable:`, rollbackError); + registry.removeSource(id); + effective.delete(id); + } + } + } + + return { + sources: [...effective.values()].map((e) => e.source), + tools: [...effective.values()].flatMap((e) => e.tools), + }; +} diff --git a/src/utils/source-mutation.ts b/src/utils/source-mutation.ts new file mode 100644 index 00000000..6ad108c3 --- /dev/null +++ b/src/utils/source-mutation.ts @@ -0,0 +1,28 @@ +/** + * Shared state for runtime source changes: which source ids came from the + * startup configuration (immutable over the API), and a lock so a TOML reload + * and an API call never interleave their per-source add/remove steps. + */ + +const fileSourceIds = new Set(); +let chain: Promise = Promise.resolve(); + +/** Record the ids the startup configuration (TOML / --dsn / demo) defines. */ +export function setFileSourceIds(ids: Iterable): void { + fileSourceIds.clear(); + for (const id of ids) { + fileSourceIds.add(id); + } +} + +/** Whether a source id is owned by the startup configuration. */ +export function isFileSource(id: string): boolean { + return fileSourceIds.has(id); +} + +/** Run `fn` after every previously queued mutation has settled. */ +export function withSourceLock(fn: () => Promise): Promise { + const run = chain.then(fn, fn); + chain = run.catch(() => undefined); + return run; +}