fix(postgres): cancel the running query when the request is aborted - #401
fix(postgres): cancel the running query when the request is aborted#401Elrendio wants to merge 2 commits into
Conversation
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 bytebase#384 on MySQL — a query outliving the client that asked for it — and this mirrors the fix accepted for it in bytebase#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 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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 <noreply@anthropic.com>
|
Pushed Caught by driving a real server over stdio against PostgreSQL and watching The unit test did not catch it because it hand-rolled |
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
src/tools/execute-sql.tsreceivesextra, which carries the MCP SDK'sAbortSignal, and uses itonly for
trackToolRequest(...). There is noAbortSignal/AbortControlleranywhere insrc/, andExecuteOptionshas no field to carry one. Cancelling a tool call therefore only stops us waitingfor the answer: the backend goes on producing a result set nobody will ever read.
This is the same class of bug as #384 ("MySQL queries can continue running on the server after
query timeout") and this change 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 inherits. PostgreSQL was simply left behind.
It matters most where it hurts most — 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
given up and moved on. We hit this running DBHub against ~24 production PostgreSQL read replicas;
abandoned agent queries contributed to replication lag of up to 10.9 h and one production incident
caused by a single query that ran for 1 h 30.
Reproduction
Added as the integration test
Query cancellation (options.signal): startSELECT pg_sleep(30),abort the signal, then watch
pg_stat_activity.expected null to be an instance of Error—the query resolved. The abort changed nothing and the backend slept the full 30 s.
query_canceled) within milliseconds,and the backend is gone from
pg_stat_activity.Fix
ExecuteOptionsgains an optionalsignal;execute_sqlpasses the request's own, read fromextra.mcpReq.signal— a tool handler receives{ sessionId, mcpReq, http }, and the AbortSignalhangs off the request context, not off the extra itself. Connectors
that ignore it are unaffected — MySQL/MariaDB/SQLite/SQL Server behaviour is untouched.
connection, since the one running the query is blocked waiting for it.
cancelBackendopens adedicated 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.
pg's ownClient.cancel(client, query)is not reusable here: it needs the internalQueryobject (
client.activeQuery === query), which the promise-basedclient.query()API neverexposes.
pg_cancel_backend(pid)over a fresh connection is the same protocol-level operation,and matches how
killQuerywas done for MySQL in fix(mysql): kill server-side query on timeout instead of leaking it into the pool #386.statement it targets, so it can instead land on the
ROLLBACKissued during cleanup and leave thesession sitting in an aborted transaction. Mirrors MySQL's
isConnectionPoisoned. Covered by atest asserting the connector still works after a cancellation.
finally, so listeners don't accumulate across the statements ofone request.
an
AbortSignaldispatchesabortexactly once, so a signal that fired while we waited for thepool would never reach a listener registered afterwards.
Cancellation stays advisory — the backend acts on it at an interrupt point, and a statement that
already finished ignores it — so
cancelBackendis best-effort and logs rather than throwing.Testing
pnpm run test:unit— 961 passed (958 onmain; +3 new).pnpm run test:integration— 336 passed (331 onmain; +5) across PostgreSQL, MySQL, MariaDB,SQL Server and SQLite containers.
pnpm run build:backend— clean.30 s, as shown above.
New coverage: signal forwarding from the tool handler (including a regression guard that a
top-level
extra.signalis not read), and four PostgreSQL cancellation integration tests (plain,inside
BEGIN READ ONLY, pool reuse after cancellation, already-aborted signal).The handler-to-connector wiring is covered only by unit tests with a hand-written
extra, which ishow the wrong path survived the first round here. An end-to-end test would need to drive a real
server over stdio against PostgreSQL; the existing JSON-RPC integration harness is HTTP + SQLite, so
that is not a small addition. I verified this path manually instead (see Reproduction).
Notes for reviewers
explain_sqlalso runs user SQL and could take the samesignal— one line — but is left out tokeep this diff to the reported bug. Happy to add it.
a natural follow-up; MySQL already has the timeout-triggered half from fix(mysql): kill server-side query on timeout instead of leaking it into the pool #386.
npx tsc --noEmitreports 132 pre-existing errors onmain; this branch reports the same 132,none in the changed files.
max_rowscoverage bug, sent as its own PR — the two are independentand share no files but one integration test file.