From dfd0f337be0d2eb3fe3b7a6facbf209fc36c2150 Mon Sep 17 00:00:00 2001 From: Elrendio Date: Sun, 9 Aug 2026 19:18:03 +0200 Subject: [PATCH 1/2] fix(postgres): cancel the running query when the request is aborted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MCP request's AbortSignal reached `execute_sql` (as `extra`) but was only ever read for telemetry. Nothing in `src/` referenced an AbortSignal, and `ExecuteOptions` had no field to carry one, so cancelling a tool call merely stopped us waiting for the answer: the backend went on producing a result set nobody would ever read. That is the same class of bug as #384 on MySQL — a query outliving the client that asked for it — and this mirrors the fix accepted for it in #386. MySQL's trigger is its own client-side timeout; here the trigger is the client cancelling. Both end the same way: kill the statement server-side over a second connection, and don't hand the connection back to the pool in a state the next borrower would inherit. It matters most exactly where it hurts most: on read replicas serving agents, where an abandoned analytical query keeps consuming IO for its full runtime and drives replication lag long after the agent has gone. Reproduction (integration test `Query cancellation (options.signal)`): start `SELECT pg_sleep(30)`, abort the signal, then watch `pg_stat_activity`. Before this change the test fails after 30,042ms with the query having *resolved* — the abort changed nothing and the backend slept the full 30s. After it, the statement rejects with SQLSTATE 57014 (query_canceled) within milliseconds and the backend is gone from `pg_stat_activity`. - `ExecuteOptions` gains an optional `signal`; `execute_sql` passes the request's own. Connectors that ignore it are unaffected. - PostgreSQL cancellation is out-of-band: the request must arrive on a different connection, since the one running the query is blocked waiting for it. `cancelBackend` opens a dedicated short-lived connection rather than borrowing from the pool — a pool saturated with the very queries being cancelled has nothing left to lend. `pg`'s own `Client.cancel` is not reusable here: it needs the internal Query object, which the promise-based query API never exposes. - The connection is destroyed instead of released after a cancellation. A CancelRequest races the statement it targets, so it can instead land on the ROLLBACK issued during cleanup and leave the session in an aborted transaction. Mirrors MySQL's `isConnectionPoisoned`. - The abort listener is removed in `finally`, so listeners don't accumulate across the statements of one request. - An already-aborted signal is refused before a connection is taken, and re-checked once one is held: an AbortSignal dispatches "abort" exactly once, so a signal that fired while we waited for the pool would never reach a listener registered afterwards. Co-Authored-By: Claude Opus 5 --- .../__tests__/postgres.integration.test.ts | 119 ++++++++++++++++++ src/connectors/interface.ts | 8 ++ src/connectors/postgres/index.ts | 103 ++++++++++++++- src/tools/__tests__/execute-sql.test.ts | 30 +++++ src/tools/execute-sql.ts | 4 + 5 files changed, 263 insertions(+), 1 deletion(-) diff --git a/src/connectors/__tests__/postgres.integration.test.ts b/src/connectors/__tests__/postgres.integration.test.ts index c632fd8..f7dcfe9 100644 --- a/src/connectors/__tests__/postgres.integration.test.ts +++ b/src/connectors/__tests__/postgres.integration.test.ts @@ -699,6 +699,125 @@ describe('PostgreSQL Connector Integration Tests', () => { }); }); + describe('Query cancellation (options.signal)', () => { + // A cancelled MCP tool call must stop the query on the *server*. Without + // that, aborting only stops us waiting for the answer while the backend + // keeps burning IO to produce a result nobody will read. + const PROBE = 'dbhub_cancel_probe'; + + /** Backends the server currently reports as running the probe query. */ + const runningProbes = async (observer: PostgresConnector): Promise => { + const result = await observer.executeSQL( + `SELECT count(*)::int AS n FROM pg_stat_activity + WHERE state = 'active' + AND query LIKE '%${PROBE}%' + AND query NOT LIKE '%pg_stat_activity%'`, + {} + ); + return result.resultSets[0].rows[0].n; + }; + + const waitFor = async (predicate: () => Promise, timeoutMs: number) => { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await predicate()) return true; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + return false; + }; + + it.each([ + ['plain execution', {} as const], + ['inside a READ ONLY transaction', { readonly: true } as const], + ])('kills the running backend when the request is aborted (%s)', async (_label, options) => { + const connector = new PostgresConnector(); + const observer = new PostgresConnector(); + try { + await connector.connect(postgresTest.connectionString); + await observer.connect(postgresTest.connectionString); + + const controller = new AbortController(); + const query = connector.executeSQL(`SELECT pg_sleep(30), '${PROBE}'`, { + ...options, + signal: controller.signal, + }); + // Capture the rejection now so polling below can't trip an unhandled one. + const settled = query.then( + () => null, + (error) => error as NodeJS.ErrnoException + ); + + expect(await waitFor(async () => (await runningProbes(observer)) === 1, 10000)).toBe(true); + + controller.abort(); + + const error = await settled; + expect(error).toBeInstanceOf(Error); + // 57014 = query_canceled: the server stopped it, we didn't just walk away. + expect((error as any).code).toBe('57014'); + + expect(await waitFor(async () => (await runningProbes(observer)) === 0, 10000)).toBe(true); + } finally { + await connector.disconnect(); + await observer.disconnect(); + } + }, 60000); + + it('leaves the connector usable after a cancellation', async () => { + // The cancelled connection is destroyed rather than returned to the pool: + // a CancelRequest races the statement it targets, so it can land on the + // cleanup ROLLBACK instead and strand the session in an aborted + // transaction that the next borrower would inherit. + const connector = new PostgresConnector(); + try { + await connector.connect(postgresTest.connectionString); + + const controller = new AbortController(); + const cancelled = connector + .executeSQL(`SELECT pg_sleep(30), '${PROBE}'`, { + readonly: true, + signal: controller.signal, + }) + .catch(() => 'rejected'); + await new Promise((resolve) => setTimeout(resolve, 500)); + controller.abort(); + expect(await cancelled).toBe('rejected'); + + const after = await connector.executeSQL('SELECT 1 AS ok', { readonly: true }); + expect(after.resultSets[0].rows[0].ok).toBe(1); + } finally { + await connector.disconnect(); + } + }, 60000); + + it('does not run the query at all when the signal is already aborted', async () => { + const connector = new PostgresConnector(); + try { + await connector.connect(postgresTest.connectionString); + + const controller = new AbortController(); + controller.abort(); + + await expect( + connector.executeSQL('SELECT 1', { signal: controller.signal }) + ).rejects.toThrow(); + } finally { + await connector.disconnect(); + } + }, 30000); + + it('runs normally when no signal is supplied', async () => { + const connector = new PostgresConnector(); + try { + await connector.connect(postgresTest.connectionString); + const result = await connector.executeSQL('SELECT 1 AS ok', {}); + expect(result.resultSets[0].rows[0].ok).toBe(1); + } finally { + await connector.disconnect(); + } + }, 30000); + }); + describe('Search Path Configuration Tests', () => { it('should use first schema in search_path as default for discovery', async () => { const connector = new PostgresConnector(); diff --git a/src/connectors/interface.ts b/src/connectors/interface.ts index 151ba2b..335400a 100644 --- a/src/connectors/interface.ts +++ b/src/connectors/interface.ts @@ -121,6 +121,14 @@ export interface ExecuteOptions { * Note: SDK-level readonly enforcement is set via ConnectorConfig.readonly */ readonly?: boolean; + /** + * Cancellation signal for the request this execution serves — the MCP + * request's own AbortSignal, so a tool call the client cancels stops the + * query on the *database* instead of leaving it running while nobody waits + * for the result. Honoured by the PostgreSQL connector; a connector that + * ignores it runs the query to completion as before. + */ + signal?: AbortSignal; } /** diff --git a/src/connectors/postgres/index.ts b/src/connectors/postgres/index.ts index d8b9139..40c6b4e 100644 --- a/src/connectors/postgres/index.ts +++ b/src/connectors/postgres/index.ts @@ -155,6 +155,10 @@ export class PostgresConnector implements Connector { private pool: pg.Pool | null = null; + // Kept so cancelBackend can open its own connection to the same server + // without borrowing from the pool + private poolConfig: pg.PoolConfig | null = null; + // Source ID is set by ConnectorManager after cloning private sourceId: string = "default"; @@ -193,6 +197,7 @@ export class PostgresConnector implements Connector { } } + this.poolConfig = poolConfig; this.pool = new Pool(poolConfig); // Test the connection @@ -206,6 +211,7 @@ export class PostgresConnector implements Connector { this.pool = null; await closeQuietly(() => pool.end()); } + this.poolConfig = null; console.error("Failed to connect to PostgreSQL database:", err); throw err; } @@ -216,6 +222,7 @@ export class PostgresConnector implements Connector { await this.pool.end(); this.pool = null; } + this.poolConfig = null; } async getSchemas(): Promise { @@ -668,13 +675,46 @@ export class PostgresConnector implements Connector { } + /** + * How long to give the out-of-band cancellation before abandoning it. + * Connecting and asking the server to signal a backend is metadata-only work + * that a healthy server answers in milliseconds; the bound exists only so + * cleanup cannot outlive the query it is cleaning up after, independently of + * the user's (possibly long, possibly unset) query timeout. + */ + private static readonly CANCEL_TIMEOUT_MS = 5000; + async executeSQL(sql: string, options: ExecuteOptions, parameters?: any[]): Promise { if (!this.pool) { throw new Error("Not connected to database"); } + const signal = options.signal; + // Nothing to run for a caller that is already gone - and no reason to take + // a connection from the pool to find that out. + signal?.throwIfAborted(); + const client = await this.pool.connect(); + // Captured up front: after the cancellation lands, reading this off the + // client races its teardown. + const backendPid = backendPidOf(client); + let wasCancelled = false; + const onAbort = () => { + wasCancelled = true; + // Fire-and-forget: an abort listener cannot be awaited, and the + // cancellation reaches the caller through the in-flight query rejecting + // with 57014 (query_canceled), not through this call. cancelBackend + // never rejects. + void this.cancelBackend(backendPid); + }; + try { + // Re-checked now that we hold a connection: an AbortSignal dispatches + // "abort" exactly once, so one that fired while we waited for the pool + // would never reach a listener registered afterwards. + signal?.throwIfAborted(); + signal?.addEventListener("abort", onAbort, { once: true }); + // Check if this is a multi-statement query const statements = splitSQLStatements(sql, "postgres"); @@ -765,9 +805,70 @@ export class PostgresConnector implements Connector { return { resultSets }; } } finally { - client.release(); + // Removed explicitly rather than left to `once`: a signal outlives a + // single executeSQL call, so listeners would otherwise pile up pointing + // at connections that have already gone back to the pool. + signal?.removeEventListener("abort", onAbort); + // release(true) destroys the connection instead of returning it to the + // pool. After a cancellation the session's state is not trustworthy: the + // CancelRequest races the statement it targets, so it can instead land + // on the ROLLBACK issued during cleanup and leave the session sitting in + // an aborted transaction that the next borrower would inherit. Same + // reasoning as the MySQL connector's isConnectionPoisoned handling. + client.release(wasCancelled); } } + + /** + * Best-effort server-side cancellation of whatever `backendPid` is running, + * for when the MCP client aborts the tool call. + * + * PostgreSQL cancellation is out-of-band by design: the request has to arrive + * on a *different* connection, because the one running the query is blocked + * waiting for its result. This opens a dedicated short-lived connection + * rather than borrowing from the pool - a pool saturated with the very + * long-running queries being cancelled has nothing left to lend, which is + * exactly the situation this runs in. That mirrors what `pg`'s own + * Client.cancel does at the protocol level; that entry point is not reusable + * here because it needs the internal Query object, which the promise-based + * query API never exposes. + * + * Cancellation is advisory: the backend acts on it only at an interrupt + * point, and a statement that has already finished ignores it. A resolved + * call is therefore not a guarantee that anything stopped. + */ + private async cancelBackend(backendPid: number | null): Promise { + if (backendPid === null || !this.poolConfig) { + return; + } + const cancelClient = new pg.Client({ + ...this.poolConfig, + connectionTimeoutMillis: PostgresConnector.CANCEL_TIMEOUT_MS, + query_timeout: PostgresConnector.CANCEL_TIMEOUT_MS, + }); + try { + await cancelClient.connect(); + await cancelClient.query("SELECT pg_cancel_backend($1)", [backendPid]); + } catch (error) { + // Unconfirmed cancellation: the statement may still be running. There is + // nothing further to do from here - the caller already sees the abort - + // but it must be visible, since it means a query outlived its request. + console.error(`Failed to cancel PostgreSQL backend ${backendPid}:`, error); + } finally { + await closeQuietly(() => cancelClient.end()); + } + } +} + +/** + * The server-side PID of the backend serving this connection. `pg` fills it in + * from the BackendKeyData message at connection time, but does not declare it + * on its published types, hence the narrow cast. Null when it is unavailable, + * which leaves the connection simply uncancellable rather than erroring. + */ +function backendPidOf(client: pg.PoolClient): number | null { + const processID = (client as unknown as { processID?: number | null }).processID; + return typeof processID === "number" ? processID : null; } // Create and register the connector diff --git a/src/tools/__tests__/execute-sql.test.ts b/src/tools/__tests__/execute-sql.test.ts index 12e7664..c7b8eae 100644 --- a/src/tools/__tests__/execute-sql.test.ts +++ b/src/tools/__tests__/execute-sql.test.ts @@ -164,6 +164,36 @@ describe('execute-sql tool', () => { }); }); + describe('cancellation', () => { + // The connector needs the request's AbortSignal to stop the query on the + // database. Without it a cancelled tool call only stops us waiting for the + // answer, and the query runs to completion server-side. + it('forwards the request AbortSignal to the connector', async () => { + vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] }); + const controller = new AbortController(); + + const handler = createExecuteSqlToolHandler('test_source'); + await handler({ sql: 'SELECT 1' }, { signal: controller.signal }); + + expect(mockConnector.executeSQL).toHaveBeenCalledWith( + 'SELECT 1', + expect.objectContaining({ signal: controller.signal }) + ); + }); + + it('passes an undefined signal when the caller provides no extra', async () => { + vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] }); + + const handler = createExecuteSqlToolHandler('test_source'); + await handler({ sql: 'SELECT 1' }, null); + + expect(mockConnector.executeSQL).toHaveBeenCalledWith( + 'SELECT 1', + expect.objectContaining({ signal: undefined }) + ); + }); + }); + describe('read-only mode enforcement', () => { // Statement-classification coverage (write keywords, comment stripping, // dialect-specific bypasses, ...) is pinned in diff --git a/src/tools/execute-sql.ts b/src/tools/execute-sql.ts index c68fcda..c6ad08c 100644 --- a/src/tools/execute-sql.ts +++ b/src/tools/execute-sql.ts @@ -63,6 +63,10 @@ export function createExecuteSqlToolHandler(sourceId?: string) { const executeOptions = { readonly: isReadOnlyPolicy(policy), maxRows: toolConfig?.max_rows, + // The MCP request's cancellation signal. Without it a cancelled tool + // call only stops us waiting for the answer: the query keeps running + // on the database, still burning the IO nobody is going to read. + signal: extra?.signal as AbortSignal | undefined, }; result = await connector.executeSQL(sql, executeOptions); From 5354636b6d10eb0c2173c05a68e4600734085ac1 Mon Sep 17 00:00:00 2001 From: Elrendio Date: Mon, 10 Aug 2026 15:17:22 +0200 Subject: [PATCH 2/2] fix(postgres): read the request signal from the MCP request context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The signal was read as `extra.signal`, which the SDK never populates: a tool handler receives `{ sessionId, mcpReq, http }`, and the request's AbortSignal hangs off `extra.mcpReq`. `options.signal` was therefore always undefined, so cancelling a tool call still left the query running on the database — the exact failure the connector change was written to prevent. Proven end to end against a real PostgreSQL, driving the server over stdio and watching pg_stat_activity: with the fix the backend disappears 201 ms after `notifications/cancelled`, and reverting only this line leaves it running past 10 s. The unit test could not catch this: it hand-rolled `{ signal }` as the handler extra, encoding the same wrong assumption as the code. It now uses the shape the SDK actually passes, plus a regression guard asserting that a top-level `extra.signal` is NOT read. Co-Authored-By: Claude Opus 5 --- src/tools/__tests__/execute-sql.test.ts | 19 ++++++++++++++++++- src/tools/execute-sql.ts | 10 ++++++---- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/tools/__tests__/execute-sql.test.ts b/src/tools/__tests__/execute-sql.test.ts index c7b8eae..39ae846 100644 --- a/src/tools/__tests__/execute-sql.test.ts +++ b/src/tools/__tests__/execute-sql.test.ts @@ -173,7 +173,8 @@ describe('execute-sql tool', () => { const controller = new AbortController(); const handler = createExecuteSqlToolHandler('test_source'); - await handler({ sql: 'SELECT 1' }, { signal: controller.signal }); + // The shape the SDK actually hands a tool handler: `{ sessionId, mcpReq, http }`. + await handler({ sql: 'SELECT 1' }, { sessionId: 's', mcpReq: { signal: controller.signal } }); expect(mockConnector.executeSQL).toHaveBeenCalledWith( 'SELECT 1', @@ -181,6 +182,22 @@ describe('execute-sql tool', () => { ); }); + // Regression guard: reading `extra.signal` instead of `extra.mcpReq.signal` type-checks and + // passes a hand-rolled `{ signal }` fixture, while silently never cancelling anything in + // production — the signal is simply always undefined. + it('does not read a signal off the handler extra itself', async () => { + vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] }); + const controller = new AbortController(); + + const handler = createExecuteSqlToolHandler('test_source'); + await handler({ sql: 'SELECT 1' }, { signal: controller.signal }); + + expect(mockConnector.executeSQL).toHaveBeenCalledWith( + 'SELECT 1', + expect.objectContaining({ signal: undefined }) + ); + }); + it('passes an undefined signal when the caller provides no extra', async () => { vi.mocked(mockConnector.executeSQL).mockResolvedValue({ resultSets: [] }); diff --git a/src/tools/execute-sql.ts b/src/tools/execute-sql.ts index c6ad08c..cc71769 100644 --- a/src/tools/execute-sql.ts +++ b/src/tools/execute-sql.ts @@ -63,10 +63,12 @@ export function createExecuteSqlToolHandler(sourceId?: string) { const executeOptions = { readonly: isReadOnlyPolicy(policy), maxRows: toolConfig?.max_rows, - // The MCP request's cancellation signal. Without it a cancelled tool - // call only stops us waiting for the answer: the query keeps running - // on the database, still burning the IO nobody is going to read. - signal: extra?.signal as AbortSignal | undefined, + // The MCP request's cancellation signal, which the SDK hangs off the + // request context (`extra.mcpReq`) rather than the handler extra + // itself. Without it a cancelled tool call only stops us waiting for + // the answer: the query keeps running on the database, still burning + // the IO nobody is going to read. + signal: extra?.mcpReq?.signal as AbortSignal | undefined, }; result = await connector.executeSQL(sql, executeOptions);