From f49ff9421480bf5a5d88f9ce6fcc6b7cdaa51e09 Mon Sep 17 00:00:00 2001 From: suibianwanwan Date: Tue, 29 Sep 2026 10:30:33 +0800 Subject: [PATCH 1/5] fix(cli): return failure exit code for SQL batches --- packages/cz-cli/src/commands/sql.ts | 6 +- .../cz-cli/test/sql-batch-exit-code.test.ts | 57 +++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 packages/cz-cli/test/sql-batch-exit-code.test.ts diff --git a/packages/cz-cli/src/commands/sql.ts b/packages/cz-cli/src/commands/sql.ts index 97aebf4f1..b237b7fb0 100644 --- a/packages/cz-cli/src/commands/sql.ts +++ b/packages/cz-cli/src/commands/sql.ts @@ -3,7 +3,7 @@ import type { Argv } from "yargs" import { readFileSync, openSync, readSync, closeSync } from "node:fs" import { analyzeSql, isReadonlySqlSetting, splitSql, JobStatus, requestRaw, getCurrentUser, type JobID, type QueryResult } from "@clickzetta/sdk" import type { GlobalArgs } from "../cli.js" -import { success, successRows, error, handledError, parseOutputArgs, renderOutput, renderErrorOutput } from "../output/index.js" +import { success, successRows, error, handledError, parseOutputArgs, renderOutput, renderErrorOutput, EXIT_BIZ_ERROR } from "../output/index.js" import { maskRows } from "../output/masking.js" import { logOperation } from "../logger.js" import { type ExecContext, classifyExecError, execSql, execSqlWithRetry, getExecContext, isQueryResult, validateIdentifier } from "./exec.js" @@ -704,6 +704,7 @@ async function handler(argv: SqlArgs): Promise { try { const r = await execSqlWithRetry(ctx, stmt, { hints: accumulatedHints, timeoutMs: argv.timeout * 1000, configStatements }) if (isQueryResult(r) && r.status === JobStatus.FAILED) { + process.exitCode = EXIT_BIZ_ERROR const line = { index: i, sql: stmt, error: { code: r.errorCode ?? "SQL_ERROR", message: await formatQueryError(r, ctx, argv.profile) }, time_ms: Date.now() - t0, ...(r.jobId ? { job_id: r.jobId } : {}) } process.stdout.write(renderOutput(line, format, batchField) + "\n") logOperation("sql", { sql: stmt, ok: false, errorCode: r.errorCode }) @@ -715,6 +716,7 @@ async function handler(argv: SqlArgs): Promise { logOperation("sql", { sql: stmt, ok: true, rows: rows.length, timeMs: Date.now() - t0 }) } } catch (err) { + process.exitCode = EXIT_BIZ_ERROR const { code, message } = classifyExecError(err) const line = { index: i, sql: stmt, error: { code, message: await formatClassifiedError({ code, message, ctx, profileName: argv.profile }) }, time_ms: Date.now() - t0 } process.stdout.write(renderOutput(line, format, batchField) + "\n") @@ -858,7 +860,7 @@ export function registerSqlCommand(cli: Argv): void { .option("header", { type: "boolean", default: true, describe: "Include column names in output. Use --no-header or -N to suppress." }) .option("N", { type: "boolean", hidden: true }) .option("limit", { type: "number", default: 100, describe: "Max rows to return (0 for unlimited)" }) - .option("batch", { alias: "B", type: "boolean", default: false, describe: "Batch mode: execute multiple semicolon-separated statements sequentially" }) + .option("batch", { alias: "B", type: "boolean", default: false, describe: "Batch mode: execute all statements sequentially; exit non-zero if any fails" }) .option("dry-run", { type: "boolean", default: false, describe: "EXPLAIN recognized readonly queries. Multiple queries require sql_split=true; writes, session commands and unknown syntax are rejected." }) .epilogue([ "Examples:", diff --git a/packages/cz-cli/test/sql-batch-exit-code.test.ts b/packages/cz-cli/test/sql-batch-exit-code.test.ts new file mode 100644 index 000000000..094e1ecc4 --- /dev/null +++ b/packages/cz-cli/test/sql-batch-exit-code.test.ts @@ -0,0 +1,57 @@ +import { beforeEach, expect, test } from "bun:test" +import { join } from "node:path" +import { onFetch, requireTestHome, sqlFailure, sqlSuccess, stubStudioContext } from "./support/cz-fixtures.js" + +const { execute } = await import("../src/execute.ts") + +beforeEach(async () => { + stubStudioContext() + await Bun.write(join(requireTestHome(), ".clickzetta", "profiles.toml"), [ + 'default_profile = "test"', + "[profiles.test]", + 'pat = "pat"', + 'service = "uat-api.clickzetta.com"', + 'instance = "inst"', + 'workspace = "ws0"', + ].join("\n")) +}) + +// Exercise the real command and SDK with only the HTTP boundary substituted. +for (const input of ["execute", "file"] as const) { + for (const failure of ["none", "job", "http"] as const) { + test(`batch ${input}: ${failure} failure preserves results and exit status`, async () => { + const submitted: string[] = [] + onFetch({ + match: (url) => url.includes("/lh/submitJob"), + respond: (_url, _method, body) => { + const query = (body as { jobDesc: { sqlJob: { query: string[] } } }).jobDesc.sqlJob.query[0] + submitted.push(query) + if (query.includes("SELECT 2") && failure === "job") { + return sqlFailure("CZLH-42000", "Statement failed") + } + if (query.includes("SELECT 2") && failure === "http") { + return new Response("Submission rejected", { status: 400 }) + } + return sqlSuccess(["value"], [[1]]) + }, + }) + + const sql = "SELECT 1; SELECT 2; SELECT 3;" + const file = join(requireTestHome(), "batch.sql") + if (input === "file") await Bun.write(file, sql) + const result = await execute("sql --batch --sync", input === "file" ? ["-f", file] : ["-e", sql]) + const rows = result.output.trim().split("\n").map((line) => JSON.parse(line)) + + expect(submitted).toHaveLength(3) + expect(rows.map((row) => row.index)).toEqual([0, 1, 2]) + expect(rows[0].rows).toEqual([[1]]) + expect(rows[2].rows).toEqual([[1]]) + expect(result.exitCode).toBe(failure === "none" ? 0 : 1) + if (failure === "none") { + expect(rows.every((row) => !row.error)).toBe(true) + return + } + expect(rows[1].error.code).toBe(failure === "job" ? "CZLH-42000" : "HTTP_400") + }) + } +} From 7574edf8c4b1cb25c53f5ea85b6962a2a9b83d58 Mon Sep 17 00:00:00 2001 From: suibianwanwan Date: Thu, 8 Oct 2026 11:14:21 +0800 Subject: [PATCH 2/5] fix(cli): fail dry-run when any EXPLAIN fails --dry-run reported per-statement errors but exited 0. Batch and dry-run failures now also record lastError so telemetry carries the error. Co-Authored-By: Claude Opus 5.5 (1M context) --- packages/cz-cli/src/commands/sql.ts | 15 ++++++++++--- .../cz-cli/test/sql-batch-exit-code.test.ts | 21 +++++++++++++++++++ 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/packages/cz-cli/src/commands/sql.ts b/packages/cz-cli/src/commands/sql.ts index b237b7fb0..e2f84b132 100644 --- a/packages/cz-cli/src/commands/sql.ts +++ b/packages/cz-cli/src/commands/sql.ts @@ -502,6 +502,13 @@ async function resolveAccountDisplayName(ctx: ExecContext) { } } +// Per-statement failures are already rendered inline, so error() would print a +// second envelope; set the exit code and telemetry error it would have set. +function markStatementFailure(message: string) { + process.exitCode = EXIT_BIZ_ERROR + ;(process as unknown as Record).lastError = message +} + async function formatQueryError(r: QueryResult, ctx: ExecContext, profileName?: string, fallback = "Query failed") { return formatBillingError({ code: r.errorCode, @@ -669,6 +676,8 @@ async function handler(argv: SqlArgs): Promise { } })) success({ statements: results, count: statements.length }, { format, rowsKey: "statements" }) + const failed = results.find((r) => r.status === "error") + if (failed) markStatementFailure(String(failed.error)) return } ctx = await getExecContext(argv) @@ -704,8 +713,8 @@ async function handler(argv: SqlArgs): Promise { try { const r = await execSqlWithRetry(ctx, stmt, { hints: accumulatedHints, timeoutMs: argv.timeout * 1000, configStatements }) if (isQueryResult(r) && r.status === JobStatus.FAILED) { - process.exitCode = EXIT_BIZ_ERROR const line = { index: i, sql: stmt, error: { code: r.errorCode ?? "SQL_ERROR", message: await formatQueryError(r, ctx, argv.profile) }, time_ms: Date.now() - t0, ...(r.jobId ? { job_id: r.jobId } : {}) } + markStatementFailure(line.error.message) process.stdout.write(renderOutput(line, format, batchField) + "\n") logOperation("sql", { sql: stmt, ok: false, errorCode: r.errorCode }) } else if (isQueryResult(r)) { @@ -716,9 +725,9 @@ async function handler(argv: SqlArgs): Promise { logOperation("sql", { sql: stmt, ok: true, rows: rows.length, timeMs: Date.now() - t0 }) } } catch (err) { - process.exitCode = EXIT_BIZ_ERROR const { code, message } = classifyExecError(err) const line = { index: i, sql: stmt, error: { code, message: await formatClassifiedError({ code, message, ctx, profileName: argv.profile }) }, time_ms: Date.now() - t0 } + markStatementFailure(line.error.message) process.stdout.write(renderOutput(line, format, batchField) + "\n") logOperation("sql", { sql: stmt, ok: false, errorCode: code }) } @@ -861,7 +870,7 @@ export function registerSqlCommand(cli: Argv): void { .option("N", { type: "boolean", hidden: true }) .option("limit", { type: "number", default: 100, describe: "Max rows to return (0 for unlimited)" }) .option("batch", { alias: "B", type: "boolean", default: false, describe: "Batch mode: execute all statements sequentially; exit non-zero if any fails" }) - .option("dry-run", { type: "boolean", default: false, describe: "EXPLAIN recognized readonly queries. Multiple queries require sql_split=true; writes, session commands and unknown syntax are rejected." }) + .option("dry-run", { type: "boolean", default: false, describe: "EXPLAIN recognized readonly queries; exit non-zero if any fails. Multiple queries require sql_split=true; writes, session commands and unknown syntax are rejected." }) .epilogue([ "Examples:", " cz-cli sql \"SELECT * FROM orders LIMIT 10\"", diff --git a/packages/cz-cli/test/sql-batch-exit-code.test.ts b/packages/cz-cli/test/sql-batch-exit-code.test.ts index 094e1ecc4..4e1314f9b 100644 --- a/packages/cz-cli/test/sql-batch-exit-code.test.ts +++ b/packages/cz-cli/test/sql-batch-exit-code.test.ts @@ -55,3 +55,24 @@ for (const input of ["execute", "file"] as const) { }) } } + +for (const failure of ["none", "job", "http"] as const) { + test(`dry-run: ${failure} failure preserves results and exit status`, async () => { + await Bun.write(join(requireTestHome(), ".clickzetta", "czcli.json"), '{"sql_split":true}') + onFetch({ + match: (url) => url.includes("/lh/submitJob"), + respond: (_url, _method, body) => { + const query = (body as { jobDesc: { sqlJob: { query: string[] } } }).jobDesc.sqlJob.query[0] + if (query.includes("SELECT 2") && failure === "job") return sqlFailure("CZLH-42000", "Statement failed") + if (query.includes("SELECT 2") && failure === "http") return new Response("Submission rejected", { status: 400 }) + return sqlSuccess(["plan"], [["ok"]]) + }, + }) + + const result = await execute("sql", ["SELECT 1; SELECT 2", "--dry-run"]) + const statements = JSON.parse(result.output).data.statements as { status: string }[] + + expect(statements.map((s) => s.status)).toEqual(failure === "none" ? ["ok", "ok"] : ["ok", "error"]) + expect(result.exitCode).toBe(failure === "none" ? 0 : 1) + }) +} From 2536ed9918538d7d29f885f2b201aa684dce018f Mon Sep 17 00:00:00 2001 From: suibianwanwan Date: Mon, 28 Sep 2026 16:14:59 +0800 Subject: [PATCH 3/5] fix(sql): cancel remote queries after subprocess timeout --- packages/clickzetta-sdk/src/abort.ts | 46 ++ packages/clickzetta-sdk/src/client.ts | 48 +- packages/clickzetta-sdk/src/index.ts | 2 + packages/clickzetta-sdk/src/sql/cancel.ts | 86 +++- packages/clickzetta-sdk/src/sql/poll.ts | 87 ++-- packages/clickzetta-sdk/src/sql/submit.ts | 50 +- packages/clickzetta-sdk/test/cancel.test.ts | 24 + .../test/fixtures/cancellation.ts | 126 +++++ packages/cz-cli/src/bootstrap/runtime.ts | 8 +- packages/cz-cli/src/commands/exec.ts | 106 +++-- packages/cz-cli/src/commands/sql-lifecycle.ts | 114 +++++ packages/cz-cli/src/commands/sql.ts | 25 +- packages/cz-cli/src/main.ts | 5 +- packages/cz-cli/src/sql/cleanup-scope.ts | 314 +++++++++++++ packages/cz-cli/src/sql/supervisor-runtime.ts | 32 ++ packages/cz-cli/test/sql-cancellation.test.ts | 171 +++++++ .../cz-cli/test/sql-cleanup-scope.test.ts | 443 ++++++++++++++++++ .../cz-cli/test/sql-timeout-options.test.ts | 45 ++ specs/sql-cancellation.md | 97 ++++ 19 files changed, 1681 insertions(+), 148 deletions(-) create mode 100644 packages/clickzetta-sdk/src/abort.ts create mode 100644 packages/clickzetta-sdk/test/cancel.test.ts create mode 100644 packages/clickzetta-sdk/test/fixtures/cancellation.ts create mode 100644 packages/cz-cli/src/commands/sql-lifecycle.ts create mode 100644 packages/cz-cli/src/sql/cleanup-scope.ts create mode 100644 packages/cz-cli/src/sql/supervisor-runtime.ts create mode 100644 packages/cz-cli/test/sql-cancellation.test.ts create mode 100644 packages/cz-cli/test/sql-cleanup-scope.test.ts create mode 100644 packages/cz-cli/test/sql-timeout-options.test.ts create mode 100644 specs/sql-cancellation.md diff --git a/packages/clickzetta-sdk/src/abort.ts b/packages/clickzetta-sdk/src/abort.ts new file mode 100644 index 000000000..6a523d022 --- /dev/null +++ b/packages/clickzetta-sdk/src/abort.ts @@ -0,0 +1,46 @@ +/** Race non-fetch boundaries (credential sources and backoff) against cancellation. */ +export function abortable(work: Promise, signal?: AbortSignal): Promise { + if (!signal) return work + return new Promise((resolve, reject) => { + const abort = () => reject(signal.reason) + if (signal.aborted) abort() + else signal.addEventListener("abort", abort, { once: true }) + work.then(resolve, reject).finally(() => signal.removeEventListener("abort", abort)) + }) +} + +export async function delay(ms: number, signal?: AbortSignal): Promise { + signal?.throwIfAborted() + return new Promise((resolve, reject) => { + const finish = () => { + signal?.removeEventListener("abort", abort) + resolve() + } + const timer = setTimeout(finish, ms) + const abort = () => { + clearTimeout(timer) + signal?.removeEventListener("abort", abort) + reject(signal?.reason) + } + signal?.addEventListener("abort", abort, { once: true }) + }) +} + +/** Explicitly own and dispose deadline timers, including across repeated requests. */ +export function abortAfter(timeoutMs: number, parent?: AbortSignal) { + const controller = new AbortController() + const abort = () => controller.abort(parent?.reason) + if (parent?.aborted) abort() + else parent?.addEventListener("abort", abort, { once: true }) + const timer = setTimeout( + () => controller.abort(new DOMException("The operation timed out", "TimeoutError")), + timeoutMs, + ) + return { + signal: controller.signal, + dispose() { + clearTimeout(timer) + parent?.removeEventListener("abort", abort) + }, + } +} diff --git a/packages/clickzetta-sdk/src/client.ts b/packages/clickzetta-sdk/src/client.ts index 49b4ce93c..d1fbe1586 100644 --- a/packages/clickzetta-sdk/src/client.ts +++ b/packages/clickzetta-sdk/src/client.ts @@ -1,3 +1,4 @@ +import { abortable, abortAfter, delay } from "./abort.js" import { ClickZettaApiError, type ApiResponse } from "./types/api.js" import type { Credential, RequestContext, TokenSource } from "./types/index.js" import { currentTraceparent } from "./traceparent.js" @@ -28,6 +29,9 @@ export interface ClientOptions { customHeaders?: Record traceparent?: string timeout?: number + /** Cancellation covers requests, credential waits and retry backoff. */ + signal?: AbortSignal + maxRetries?: number /** Non-auth metadata some request bodies embed — see {@link RequestContext}. */ context?: RequestContext } @@ -41,10 +45,6 @@ export function retryDelayMs(attempt: number): number { return base + Math.random() * 500 } -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)) -} - /** * Generate a request id matching the Python connector format * (`pysdk-v{version}-{uuid12}`, client.py:292). The server uses this for @@ -72,12 +72,12 @@ function buildHeaders(opts: ClientOptions, credential: Credential): Record( parseWrapper: boolean, ): Promise { const url = `${opts.baseUrl}${path}` - let credential = await opts.tokens.get() + opts.signal?.throwIfAborted() + let credential = await abortable(opts.tokens.get(), opts.signal) let headers = buildHeaders(opts, credential) // Credential a rotation just produced, to be used verbatim by the next attempt. // `TokenSource.rotate` is contracted to RETURN the replacement, not to make @@ -123,13 +124,15 @@ async function doRequest( let authExhausted = false let lastError: Error | undefined - for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) { + for (let attempt = 0; attempt <= (opts.maxRetries ?? MAX_RETRIES); attempt++) { + const deadline = abortAfter(opts.timeout ?? DEFAULT_TIMEOUT_MS, opts.signal) try { if (attempt > 0) { // Prefer the rotated credential; otherwise re-resolve, because a retry // after a multi-second backoff must not resend one that expired while // we waited. `get()` is cache-backed, so re-resolving costs nothing. - credential = rotatedCredential ?? await opts.tokens.get() + opts.signal?.throwIfAborted() + credential = rotatedCredential ?? (await abortable(opts.tokens.get(), opts.signal)) rotatedCredential = undefined headers = buildHeaders(opts, credential) } @@ -137,7 +140,7 @@ async function doRequest( method, headers, body: body !== undefined ? JSON.stringify(body) : undefined, - signal: AbortSignal.timeout(opts.timeout ?? DEFAULT_TIMEOUT_MS), + signal: deadline.signal, }) const text = await resp.text() if (!resp.ok) { @@ -150,9 +153,10 @@ async function doRequest( if (resp.status === AUTH_EXPIRED_STATUS) { // Rotation is offered once; a source that cannot rotate (or a second // rejection) makes this 401 the final answer for this identity. - const fresh = rotated || attempt >= MAX_RETRIES - ? undefined - : await opts.tokens.rotate(credential) + const fresh = + rotated || attempt >= (opts.maxRetries ?? MAX_RETRIES) + ? undefined + : await abortable(opts.tokens.rotate(credential), opts.signal) if (!fresh) { authExhausted = true throw apiErr @@ -169,16 +173,22 @@ async function doRequest( throw new ClickZettaApiError("PARSE_ERROR", `Invalid JSON response: ${text.slice(0, 200)}`, 0) } } catch (err) { + opts.signal?.throwIfAborted() lastError = err instanceof Error ? err : new Error(String(err)) - if (err instanceof ClickZettaApiError && (NON_RETRYABLE_STATUS.has(err.statusCode ?? 0) || err.code === "PARSE_ERROR")) { + if ( + err instanceof ClickZettaApiError && + (NON_RETRYABLE_STATUS.has(err.statusCode ?? 0) || err.code === "PARSE_ERROR") + ) { throw err } if (authExhausted) throw err if (TERMINAL_ERROR_CODES.has(String((err as { code?: unknown }).code ?? ""))) throw err - if (attempt < MAX_RETRIES) { - await sleep(retryDelayMs(attempt)) + if (attempt < (opts.maxRetries ?? MAX_RETRIES)) { + await delay(retryDelayMs(attempt), opts.signal) continue } + } finally { + deadline.dispose() } } // Signal to callers: parseWrapper is unused on the error path but @@ -193,7 +203,7 @@ export async function request( body?: unknown, method: string = "POST", ): Promise> { - return doRequest>(options, path, body, method, true) + return abortable(doRequest>(options, path, body, method, true), options.signal) } /** @@ -207,5 +217,5 @@ export async function requestRaw( body?: unknown, method: string = "POST", ): Promise { - return doRequest(options, path, body, method, false) + return abortable(doRequest(options, path, body, method, false), options.signal) } diff --git a/packages/clickzetta-sdk/src/index.ts b/packages/clickzetta-sdk/src/index.ts index f7039f77c..e5880cd3d 100644 --- a/packages/clickzetta-sdk/src/index.ts +++ b/packages/clickzetta-sdk/src/index.ts @@ -47,3 +47,5 @@ export const threadsafety = 2 /** DB-API 2.0 parameter style (dbapi.py:29). */ export const paramstyle = "qmark" export { czStruct } from "./sql/converter.js" + +export { abortable, abortAfter } from "./abort.js" diff --git a/packages/clickzetta-sdk/src/sql/cancel.ts b/packages/clickzetta-sdk/src/sql/cancel.ts index 076c315e1..cfa3c4436 100644 --- a/packages/clickzetta-sdk/src/sql/cancel.ts +++ b/packages/clickzetta-sdk/src/sql/cancel.ts @@ -1,21 +1,81 @@ -import { request, type ClientOptions } from "../client.js" +import { requestRaw, type ClientOptions } from "../client.js" +import { ClickZettaApiError } from "../types/api.js" +import { abortAfter, delay } from "../abort.js" +import { isRetryableErrorCode } from "./errors.js" +import { getJobResultRaw } from "./job-info.js" import type { JobID } from "./types.js" -export async function cancelJob( - opts: ClientOptions, - jobId: JobID, -): Promise { - const body = { +export async function cancelJob(opts: ClientOptions, jobId: JobID): Promise { + const response = await requestRaw(opts, "/lh/cancelJob", { account: { user_id: 0 }, - job_id: { - id: jobId.id, - workspace: jobId.workspace, - instance_id: jobId.instanceId, - }, + job_id: { id: jobId.id, workspace: jobId.workspace, instance_id: jobId.instanceId }, user_agent: "", force: false, + }) + // Coordinator protobuf JSON uses respStatus; some gateways preserve snake_case. + // Proto3 omits empty fields, so {} is a valid successful CancelJobResponse. + if (!response || typeof response !== "object" || Array.isArray(response)) { + throw new ClickZettaApiError("INVALID_CANCEL_RESPONSE", "Invalid cancellation response") + } + const raw = response as Record + const value = raw.respStatus ?? raw.resp_status + if (value !== undefined && (!value || typeof value !== "object" || Array.isArray(value))) { + throw new ClickZettaApiError("INVALID_CANCEL_RESPONSE", "Invalid cancellation status") + } + if (Object.keys(raw).length > 0 && value === undefined && raw.code === undefined) { + throw new ClickZettaApiError("INVALID_CANCEL_RESPONSE", "Missing cancellation status") + } + const status = value as Record | undefined + const code = status?.errorCode ?? status?.error_code + const message = status?.errorMsg ?? status?.error_msg + if (code || message) throw new ClickZettaApiError(String(code || "CANCEL_FAILED"), String(message || code)) + if (raw.code !== undefined && ![0, "0", 200, "200", "SUCCESS"].includes(raw.code as string | number)) { + throw new ClickZettaApiError(String(raw.code), String(raw.message ?? raw.msg ?? "Cancellation rejected")) } + return response +} + +export type CancellationResult = { confirmed: true; state: string } | { confirmed: false; reason: string } - const resp = await request(opts, "/lh/cancelJob", body) - return resp +/** Independent total budget, including credential resolution, requests and retries. */ +export async function cancelJobAndWait( + opts: ClientOptions, + jobId: JobID, + timeoutMs = 5000, +): Promise { + const deadline = abortAfter(timeoutMs) + const client = { ...opts, signal: deadline.signal, maxRetries: 0 } + try { + let reason = "Cancellation was not confirmed before the cleanup deadline" + while (!client.signal.aborted) { + // Repeat cancellation to cover a submit that becomes visible after the first cancel. + // A missing job is not proof of termination: submission may still be in flight. + await cancelJob(client, jobId).catch((error: unknown) => { + reason = + error instanceof ClickZettaApiError ? `Cancellation rejected (${error.code})` : "Cancellation request failed" + }) + const raw = await getJobResultRaw(client, jobId).catch(() => undefined) + if (raw && typeof raw === "object" && "status" in raw) { + const response = raw as { + status?: { state?: string; errorCode?: string } + respStatus?: { errorCode?: string } + resp_status?: { error_code?: string } + } + const status = response.status + if ( + status?.state && + !response.respStatus?.errorCode && + !response.resp_status?.error_code && + !isRetryableErrorCode(status?.errorCode) && + ["SUCCEED", "FAILED", "CANCELLED"].includes(status?.state ?? "") + ) { + return { confirmed: true, state: status.state } + } + } + await delay(250, client.signal).catch(() => {}) + } + return { confirmed: false, reason } + } finally { + deadline.dispose() + } } diff --git a/packages/clickzetta-sdk/src/sql/poll.ts b/packages/clickzetta-sdk/src/sql/poll.ts index 4b0f13c36..734785dfb 100644 --- a/packages/clickzetta-sdk/src/sql/poll.ts +++ b/packages/clickzetta-sdk/src/sql/poll.ts @@ -1,11 +1,7 @@ +import { abortable, delay } from "../abort.js" import { requestRaw, type ClientOptions } from "../client.js" import { JobStatus, type JobID, type QueryResult, type ColumnSchema } from "./types.js" -import { - isFatalErrorCode, - isRetryableErrorCode, - isRetryableMessage, - shouldResubmitWithNewJobId, -} from "./errors.js" +import { isFatalErrorCode, isRetryableErrorCode, isRetryableMessage, shouldResubmitWithNewJobId } from "./errors.js" import { toClickZettaError, OperationalError } from "../types/errors.js" import { decodeArrowPayload, deduplicateColumns, fetchArrowFromUrls } from "./arrow.js" import { normalizeServiceEndpoint } from "./submit.js" @@ -144,7 +140,6 @@ function splitSingle(row: string, delimiterIndex: number, columnCount: number): return result } - /** * Parse base64-encoded TEXT data chunks into rows of column values. * Each element in dataList is a base64-encoded string containing @@ -193,7 +188,8 @@ async function fetchTextFromUrls(urls: string[], columnCount: number): Promise<( function toJobStatus(state: string): JobStatus { switch (state) { - case "SUBMITTED": return JobStatus.SUBMITTED + case "SUBMITTED": + return JobStatus.SUBMITTED case "RUNNING": case "QUEUEING": case "SETUP": @@ -202,15 +198,16 @@ function toJobStatus(state: string): JobStatus { case "SUCCEEDED": case "SUCCEED": return JobStatus.SUCCEEDED - case "FAILED": return JobStatus.FAILED + case "FAILED": + return JobStatus.FAILED case "CANCELLED": case "CANCELLING": return JobStatus.CANCELLED - default: return JobStatus.UNKNOWN + default: + return JobStatus.UNKNOWN } } - // --- Type coercion --- /** @@ -282,18 +279,34 @@ export function coerceValue(value: string | null, typeCategory: string, timezone // JSON if (upper === "JSON") { - try { return JSON.parse(value) } catch { return value } + try { + return JSON.parse(value) + } catch { + return value + } } // Complex containers (JSON-encoded strings in TEXT format) if (upper === "MAP" && value) { - try { return JSON.parse(value) } catch { return value } + try { + return JSON.parse(value) + } catch { + return value + } } if ((upper.startsWith("ARRAY") || upper === "LIST") && value) { - try { return JSON.parse(value) } catch { return value } + try { + return JSON.parse(value) + } catch { + return value + } } if ((upper === "STRUCT" || upper.startsWith("STRUCT<") || upper.startsWith("ROW(")) && value) { - try { return JSON.parse(value) } catch { return value } + try { + return JSON.parse(value) + } catch { + return value + } } // String-like: CHAR / CHAR(n) / VARCHAR / VARCHAR(n) / STRING @@ -376,16 +389,11 @@ function parseResultSet( return { columns, rows: [], isAsync: false, timeZone, format } } - /** * Parse a raw LH job response (from either submitJob HYBRID or getJob) * into a QueryResult. Handles embedded TEXT data and presigned URL fetching. */ -export async function parseJobResponse( - raw: LhJobResponse, - jobId: JobID, - timezone?: string, -): Promise { +export async function parseJobResponse(raw: LhJobResponse, jobId: JobID, timezone?: string): Promise { const state = raw.status?.state ?? "UNKNOWN" const status = toJobStatus(state) @@ -440,11 +448,15 @@ export async function parseJobResponse( /** * Poll /lh/getJob until the job reaches a terminal state, then parse results. */ -export async function pollJobResult( +export function pollJobResult( opts: ClientOptions, jobId: JobID, params: PollJobResultParams = {}, ): Promise { + return abortable(pollUntilTerminal(opts, jobId, params), opts.signal) +} + +async function pollUntilTerminal(opts: ClientOptions, jobId: JobID, params: PollJobResultParams): Promise { const startTime = Date.now() const { jobTimeoutMs, timezone, maxRetries = Infinity, resubmitFn } = params const serviceInfo = normalizeServiceEndpoint(opts.context?.service ?? opts.baseUrl) @@ -468,6 +480,7 @@ export async function pollJobResult( let retryCount = 0 while (true) { + opts.signal?.throwIfAborted() if (maxRetries > 0 && retryCount >= maxRetries) { throw new OperationalError(`Job ${jobId.id} exceeded max retries (${maxRetries})`, { jobId: jobId.id }) } @@ -476,17 +489,23 @@ export async function pollJobResult( try { const { cancelJob } = await import("./cancel.js") await cancelJob(opts, jobId) - } catch { /* best-effort cancel */ } + } catch { + /* best-effort cancel */ + } throw new OperationalError(`Job ${jobId.id} timed out after ${jobTimeoutMs}ms`, { jobId: jobId.id }) } - const raw = await requestRaw({ - ...opts, - customHeaders: { - ...opts.customHeaders, - instanceId: String(jobId.instanceId), + const raw = await requestRaw( + { + ...opts, + customHeaders: { + ...opts.customHeaders, + instanceId: String(jobId.instanceId), + }, }, - }, "/lh/getJob", requestBody) + "/lh/getJob", + requestBody, + ) const state = raw?.status?.state ?? "UNKNOWN" const errorCode = raw?.status?.errorCode || raw?.respStatus?.errorCode || undefined const errorMessage = raw?.status?.errorMessage || raw?.status?.message || raw?.respStatus?.errorMsg || undefined @@ -512,16 +531,14 @@ export async function pollJobResult( // Other retryable lh_codes (60007 / 60022 / 60023) → keep polling same job if (isRetryableErrorCode(errorCode)) { retryCount++ - await new Promise((resolve) => setTimeout(resolve, sleepMs)) + await delay(sleepMs, opts.signal) sleepMs = nextSleepMs(sleepMs) continue } if (!errorCode && isRetryableMessage(errorMessage)) { const isNoPerm = - !!errorMessage && - errorMessage.includes("NoPermission: User ") && - errorMessage.endsWith(" is not found") + !!errorMessage && errorMessage.includes("NoPermission: User ") && errorMessage.endsWith(" is not found") if (isNoPerm) { noPermissionTries++ if (noPermissionTries >= NO_PERMISSION_MAX_TRIES) { @@ -532,7 +549,7 @@ export async function pollJobResult( } } retryCount++ - await new Promise((resolve) => setTimeout(resolve, sleepMs)) + await delay(sleepMs, opts.signal) sleepMs = nextSleepMs(sleepMs) continue } @@ -548,7 +565,7 @@ export async function pollJobResult( } retryCount++ - await new Promise((resolve) => setTimeout(resolve, sleepMs)) + await delay(sleepMs, opts.signal) sleepMs = nextSleepMs(sleepMs) } } diff --git a/packages/clickzetta-sdk/src/sql/submit.ts b/packages/clickzetta-sdk/src/sql/submit.ts index a1eb2cfb5..f51c4a0c4 100644 --- a/packages/clickzetta-sdk/src/sql/submit.ts +++ b/packages/clickzetta-sdk/src/sql/submit.ts @@ -1,3 +1,4 @@ +import { abortable, delay } from "../abort.js" import { requestRaw, type ClientOptions } from "../client.js" import { ClickZettaApiError } from "../types/api.js" import { lh_code } from "./errors.js" @@ -77,10 +78,6 @@ export function normalizeServiceEndpoint(value: string | undefined): { host: str } } -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)) -} - function nextRetrySleepMs(current: number): number { return current < 3000 ? current * 2 : current } @@ -99,10 +96,11 @@ function shouldRetrySubmitError(error: unknown): boolean { return ![400, 401, 403, 404, 409, 422].includes(error.statusCode ?? 0) } -export async function submitJob( - opts: ClientOptions, - params: SubmitJobParams, -): Promise { +export function submitJob(opts: ClientOptions, params: SubmitJobParams): Promise { + return abortable(submitJobRequest(opts, params), opts.signal) +} + +async function submitJobRequest(opts: ClientOptions, params: SubmitJobParams): Promise { const { sql, workspace, @@ -143,8 +141,7 @@ export async function submitJob( } } - const hybridPollingTimeout = - pollingTimeout !== undefined ? pollingTimeout : asynchronous ? 0 : 30 + const hybridPollingTimeout = pollingTimeout !== undefined ? pollingTimeout : asynchronous ? 0 : 30 const jobDesc: Record = { virtualCluster: vcluster, @@ -187,9 +184,10 @@ export async function submitJob( if (serviceInfo.endpoint) { jobDesc.jdbcDomain = serviceInfo.endpoint } - const accessToken = typeof resolvedContextJson.configs === "object" && resolvedContextJson.configs - ? (resolvedContextJson.configs as Record).access_token - : undefined + const accessToken = + typeof resolvedContextJson.configs === "object" && resolvedContextJson.configs + ? (resolvedContextJson.configs as Record).access_token + : undefined if (typeof accessToken === "string" && accessToken) { jobDesc.account = { accessToken } } @@ -200,25 +198,31 @@ export async function submitJob( let lastError: unknown = new Error(`submitJob failed for ${jobId.id}`) for (let attempt = 1; attempt <= maxRetries; attempt++) { try { - const resp = await requestRaw({ - ...opts, - traceparent, - customHeaders: { - ...opts.customHeaders, - ...(instanceName ? { instanceName } : {}), - jobId: jobId.id, + opts.signal?.throwIfAborted() + const resp = await requestRaw( + { + ...opts, + traceparent, + customHeaders: { + ...opts.customHeaders, + ...(instanceName ? { instanceName } : {}), + jobId: jobId.id, + }, }, - }, "/lh/submitJob", body) + "/lh/submitJob", + body, + ) if (submitErrorCode(resp) === lh_code.JOB_NOT_SUBMITTED && attempt < maxRetries) { - await sleep(sleepMs) + await delay(sleepMs, opts.signal) sleepMs = nextRetrySleepMs(sleepMs) continue } return resp } catch (error) { + opts.signal?.throwIfAborted() lastError = error if (attempt >= maxRetries || !shouldRetrySubmitError(error)) throw error - await sleep(sleepMs) + await delay(sleepMs, opts.signal) sleepMs = nextRetrySleepMs(sleepMs) } } diff --git a/packages/clickzetta-sdk/test/cancel.test.ts b/packages/clickzetta-sdk/test/cancel.test.ts new file mode 100644 index 000000000..b23a99700 --- /dev/null +++ b/packages/clickzetta-sdk/test/cancel.test.ts @@ -0,0 +1,24 @@ +import { expect, test } from "bun:test" + +test("SQL cancellation against real HTTP boundaries", async () => { + // Other SDK suites replace global fetch. Keep real socket/abort semantics + // isolated rather than replacing fetch again or depending on test order. + const child = Bun.spawn([process.execPath, "test", "./test/fixtures/cancellation.ts"], { + cwd: new URL("..", import.meta.url).pathname, + stdout: "pipe", + stderr: "pipe", + }) + const deadline = setTimeout(() => child.kill("SIGKILL"), 10000) + try { + const [code, stdout, stderr] = await Promise.all([ + child.exited, + new Response(child.stdout).text(), + new Response(child.stderr).text(), + ]) + expect({ code, failures: code === 0 ? "" : stdout + stderr }).toEqual({ code: 0, failures: "" }) + } finally { + clearTimeout(deadline) + child.kill("SIGKILL") + await child.exited + } +}, 15000) diff --git a/packages/clickzetta-sdk/test/fixtures/cancellation.ts b/packages/clickzetta-sdk/test/fixtures/cancellation.ts new file mode 100644 index 000000000..44f080c21 --- /dev/null +++ b/packages/clickzetta-sdk/test/fixtures/cancellation.ts @@ -0,0 +1,126 @@ +import { expect, test } from "bun:test" +import { anonymous, cancelJob, cancelJobAndWait, requestRaw, submitJob } from "../../src/index.js" + +const job = { id: "owned-job", workspace: "workspace", instanceId: 1 } + +for (const body of [ + { respStatus: { errorCode: "DENIED", errorMsg: "permission denied" } }, + { resp_status: { error_code: "DENIED", error_msg: "permission denied" } }, + { code: "DENIED", message: "permission denied" }, + null, +]) { + test(`rejects cancellation business errors: ${JSON.stringify(body)}`, async () => { + const server = Bun.serve({ port: 0, fetch: () => Response.json(body) }) + try { + await expect(cancelJob({ baseUrl: server.url.origin, tokens: anonymous() }, job)).rejects.toThrow() + } finally { + await server.stop(true) + } + }) +} + +test("retries cancellation until a late submission becomes terminal", async () => { + let attempts = 0 + const server = Bun.serve({ + port: 0, + fetch(request) { + if (request.url.endsWith("cancelJob")) { + attempts++ + return Response.json(attempts === 1 ? { respStatus: { errorCode: "CZLH-60005" } } : {}) + } + return Response.json({ + status: + attempts < 3 ? { state: "FAILED", errorCode: "CZLH-60005" } : { state: "CANCELLED", errorCode: "CZLH-60006" }, + }) + }, + }) + try { + const result = await cancelJobAndWait( + { baseUrl: server.url.origin, tokens: anonymous(), signal: AbortSignal.abort() }, + job, + 1000, + ) + expect(result).toEqual({ confirmed: true, state: "CANCELLED" }) + expect(attempts).toBe(3) + } finally { + await server.stop(true) + } +}) + +test("HTTP success is not proof of queue release", async () => { + const server = Bun.serve({ port: 0, fetch: () => Response.json({ status: { state: "RUNNING" } }) }) + try { + const result = await cancelJobAndWait({ baseUrl: server.url.origin, tokens: anonymous() }, job, 120) + expect(result.confirmed).toBe(false) + } finally { + await server.stop(true) + } +}) + +test("cleanup deadline covers credential resolution", async () => { + const start = Date.now() + const result = await cancelJobAndWait( + { + baseUrl: "http://unused.invalid", + tokens: { + get: () => new Promise(() => {}), + rotate: async () => undefined, + }, + }, + job, + 50, + ) + expect(result.confirmed).toBe(false) + expect(Date.now() - start).toBeLessThan(1000) +}) + +test("abort prevents HTTP retry after a pending request", async () => { + let requests = 0 + const server = Bun.serve({ + port: 0, + fetch() { + requests++ + return new Promise(() => {}) + }, + }) + try { + await expect( + requestRaw({ baseUrl: server.url.origin, tokens: anonymous(), signal: AbortSignal.timeout(50) }, "/pending"), + ).rejects.toThrow() + await Bun.sleep(100) + expect(requests).toBe(1) + } finally { + await server.stop(true) + } +}) + +test("abort prevents business-level submit retries", async () => { + let requests = 0 + const server = Bun.serve({ + port: 0, + fetch() { + requests++ + return Response.json({ respStatus: { errorCode: "CZLH-60023" } }) + }, + }) + try { + await expect( + submitJob( + { baseUrl: server.url.origin, tokens: anonymous(), signal: AbortSignal.timeout(50) }, + { + sql: "select 1", + workspace: "workspace", + schema: "public", + vcluster: "vc", + instanceName: "instance", + instanceId: 1, + jobId: job, + maxRetries: 3, + }, + ), + ).rejects.toThrow() + expect(requests).toBe(1) + } finally { + await server.stop(true) + } +}) diff --git a/packages/cz-cli/src/bootstrap/runtime.ts b/packages/cz-cli/src/bootstrap/runtime.ts index c5aea9699..3466e8152 100644 --- a/packages/cz-cli/src/bootstrap/runtime.ts +++ b/packages/cz-cli/src/bootstrap/runtime.ts @@ -55,6 +55,12 @@ export function applyServeLogFlags(flags: ServeLogFlags): void { } export async function main(args: string[], agentRuntime = false): Promise { + if (!agentRuntime) return runRuntime(args, false) + const { withSqlSupervisor } = await import("../sql/supervisor-runtime.js") + return withSqlSupervisor(() => runRuntime(args, true)) +} + +async function runRuntime(args: string[], agentRuntime: boolean): Promise { // cz_change: apply the base opencode env injection (kill upstream auto-updater, // disable repo-local project config, telemetry defaults) at the very top of main() // — before opencode or the TUI server Worker reads any flag. All injection is @@ -483,6 +489,6 @@ export async function main(args: string[], agentRuntime = false): Promise[0], jobId, timezone) + } } else { - result = await parseJobResponse(submitResp as Parameters[0], jobId, timezone) + result = await pollJobResult(client, jobId, { jobTimeoutMs: opts?.timeoutMs, timezone }) } - } else { - result = await pollJobResult(ctx.clientOpts, jobId, { jobTimeoutMs: opts?.timeoutMs, timezone }) - } - // Volume SQL (PUT/GET): process file transfers after getting the job result - if (isVolumeSql(normalizedSql) && result.status === JobStatus.SUCCEEDED) { - return processVolumeSql( - { clientOpts: ctx.clientOpts, workspace: ctx.config.workspace, instanceId: ctx.instanceId() }, - jobId, - result, - normalizedSql, - ) - } + disposition = "terminal" - return result + // Volume SQL (PUT/GET): process file transfers after getting the job result + if (isVolumeSql(normalizedSql) && result.status === JobStatus.SUCCEEDED) { + return processVolumeSql( + { clientOpts: ctx.clientOpts, workspace: ctx.config.workspace, instanceId: ctx.instanceId() }, + jobId, + result, + normalizedSql, + ) + } + + return result + } catch (error) { + if (lease.signal.aborted) { + const failure = new Error(`Job ${jobId.id} interrupted or timed out`) + Object.assign(failure, { jobId: jobId.id }) + throw failure + } + throw error + } finally { + await lease.finish(disposition) + } } function isAuthError(err: unknown): boolean { diff --git a/packages/cz-cli/src/commands/sql-lifecycle.ts b/packages/cz-cli/src/commands/sql-lifecycle.ts new file mode 100644 index 000000000..824a308a6 --- /dev/null +++ b/packages/cz-cli/src/commands/sql-lifecycle.ts @@ -0,0 +1,114 @@ +import { abortAfter, cancelJobAndWait, type ClientOptions, type JobID } from "@clickzetta/sdk" +import { registerSqlCleanup } from "../sql/cleanup-scope.js" +import { parseOutputArgs, renderErrorOutput } from "../output/index.js" + +const active = new Map }>() +let shutdown: Promise | undefined + +export function hasActiveSqlJobs() { + return active.size > 0 || shutdown !== undefined +} + +/** Own the remote job from before registration/submission until terminal or handoff. */ +export async function trackSqlJob(opts: ClientOptions, job: JobID, timeoutMs?: number) { + if (shutdown) throw new Error("SQL execution is stopping") + const controller = new AbortController() + const deadline = timeoutMs === undefined ? undefined : abortAfter(timeoutMs, controller.signal) + const signal = deadline?.signal ?? controller.signal + let supervisor: Awaited> + let completion: Promise | undefined + const lease = { + abort: () => controller.abort(new Error("SQL execution interrupted")), + finish(state: "terminal" | "detached" | "cancel") { + if (completion) return completion + completion = (async () => { + try { + if (state === "cancel") { + const result = await cancelJobAndWait(opts, job, 1500) + if (!result.confirmed) { + process.stderr.write( + `SQL job ${job.id}: cancellation unconfirmed; parent cleanup or server timeout must recover it.\n`, + ) + return + } + } + // A detached job is delivered only after its supervisor acknowledges handoff. + await supervisor?.release().catch(async () => { + if (state === "detached") { + const result = await cancelJobAndWait(opts, job, 1500) + if (!result.confirmed) { + process.stderr.write(`SQL job ${job.id}: cancellation after failed async handoff is unconfirmed.\n`) + } + throw Object.assign(new Error(`SQL job ${job.id}: async handoff acknowledgement failed`), { + jobId: job.id, + }) + } + // The parent can re-confirm a terminal job if its acknowledgement was lost. + process.stderr.write(`SQL job ${job.id}: cleanup acknowledgement failed.\n`) + }) + } finally { + deadline?.dispose() + supervisor?.abandon() + active.delete(job.id) + if (active.size === 0 && !shutdown) { + process.removeListener("SIGINT", interrupt) + process.removeListener("SIGTERM", terminate) + } + } + })() + return completion + }, + } + if (active.size === 0) { + process.on("SIGINT", interrupt) + process.on("SIGTERM", terminate) + } + active.set(job.id, lease) + try { + supervisor = await registerSqlCleanup(opts, job, signal, timeoutMs ?? 300_000, lease.abort) + signal.throwIfAborted() + return { + signal, + timeoutMs, + async finish(state: "terminal" | "detached" | "cancel") { + await lease.finish(state) + // Let the signal owner emit the interruption and choose the exit code. + if (shutdown) await shutdown + }, + } + } catch (error) { + await lease.finish("terminal").catch(() => {}) + throw error + } +} + +function interrupt() { + stop("SIGINT", 130) +} +function terminate() { + stop("SIGTERM", 143) +} + +function stop(signal: string, exitCode: number) { + if (shutdown) return + // Install shutdown before aborting: callbacks cannot admit another SQL job. + shutdown = Promise.resolve().then(async () => { + const jobs = [...active.entries()] + jobs.forEach(([, lease]) => lease.abort()) + const deadline = setTimeout(() => process.exit(exitCode), 2000) + await Promise.allSettled(jobs.map(([, lease]) => lease.finish("cancel"))) + const output = parseOutputArgs(process.argv.slice(2)) + process.stdout.write( + renderErrorOutput( + { + error: { code: "ABORTED", message: `Execution interrupted by ${signal}.` }, + job_ids: jobs.map(([id]) => id), + }, + output.format, + output.field, + ) + "\n", + ) + clearTimeout(deadline) + process.exit(exitCode) + }) +} diff --git a/packages/cz-cli/src/commands/sql.ts b/packages/cz-cli/src/commands/sql.ts index e2f84b132..dabd3a548 100644 --- a/packages/cz-cli/src/commands/sql.ts +++ b/packages/cz-cli/src/commands/sql.ts @@ -3,7 +3,7 @@ import type { Argv } from "yargs" import { readFileSync, openSync, readSync, closeSync } from "node:fs" import { analyzeSql, isReadonlySqlSetting, splitSql, JobStatus, requestRaw, getCurrentUser, type JobID, type QueryResult } from "@clickzetta/sdk" import type { GlobalArgs } from "../cli.js" -import { success, successRows, error, handledError, parseOutputArgs, renderOutput, renderErrorOutput, EXIT_BIZ_ERROR } from "../output/index.js" +import { success, successRows, error, handledError, parseOutputArgs, renderOutput, EXIT_BIZ_ERROR } from "../output/index.js" import { maskRows } from "../output/masking.js" import { logOperation } from "../logger.js" import { type ExecContext, classifyExecError, execSql, execSqlWithRetry, getExecContext, isQueryResult, validateIdentifier } from "./exec.js" @@ -602,6 +602,10 @@ async function emitResult( async function handler(argv: SqlArgs): Promise { const format = argv.format + if (!Number.isFinite(argv.timeout) || argv.timeout <= 0) { + error("USAGE_ERROR", "--timeout must be a positive, finite number of seconds (0 is not supported).", { format, exitCode: 2 }) + return + } if (argv["job-profile"]) { const ctx = await getExecContext(argv) @@ -635,19 +639,8 @@ async function handler(argv: SqlArgs): Promise { sql = applyVariables(sql, parseKvPairs(argv.variable)) } const hints = argv.set ? parseKvPairs(argv.set) : undefined - let currentJobId: string | undefined let ctx: ExecContext | undefined - const sigintHandler = () => { - const payload: Record = { error: { code: "ABORTED", message: "Execution interrupted by user." } } - if (currentJobId) payload.job_id = currentJobId - // renderErrorOutput, like every other failure: under a row format this is - // `ERROR ABORTED: …` rather than a JSON blob (see its docstring). - process.stdout.write(renderErrorOutput(payload, format, parseOutputArgs(process.argv.slice(2)).field) + "\n") - process.exit(130) - } - process.on("SIGINT", sigintHandler) - try { const splitEnabled = await isSplitEnabled() // Check the entire substituted input before USE validation, EXPLAIN, async, @@ -745,18 +738,16 @@ async function handler(argv: SqlArgs): Promise { return } } else { - await executeSingle(ctx, stmt, argv, accumulatedHints, configStatements, (id) => { currentJobId = id }) + await executeSingle(ctx, stmt, argv, accumulatedHints, configStatements, undefined) } } } else { - await executeSingle(ctx, statements[0], argv, hints ?? {}, undefined, (id) => { currentJobId = id }, { verbatim: !splitEnabled }) + await executeSingle(ctx, statements[0], argv, hints ?? {}, undefined, undefined, { verbatim: !splitEnabled }) } } catch (err) { const { code, message, aiMessage, jobId } = classifyExecError(err) logOperation("sql", { sql, ok: false, errorCode: code }) error(code, await formatClassifiedError({ code, message, ctx, profileName: argv.profile }), { format, debug: argv.debug, ...(aiMessage && { aiMessage }), ...(jobId && { extra: { job_id: jobId } }) }) - } finally { - process.removeListener("SIGINT", sigintHandler) } } @@ -862,7 +853,7 @@ export function registerSqlCommand(cli: Argv): void { .option("stdin", { type: "boolean", default: false, describe: "Read SQL from stdin" }) .option("sync", { type: "boolean", default: true, describe: "Wait for query result before returning (default). Use --no-sync or --async for large queries that may take a long time." }) .option("async", { type: "boolean", default: false, describe: "Return job_id immediately without waiting for results. Use for large/long-running queries." }) - .option("timeout", { type: "number", default: 300, describe: "Job timeout in seconds (default: 300)" }) + .option("timeout", { type: "number", default: 300, describe: "Positive job timeout in seconds (default: 300; 0 is not supported)" }) .option("variable", { type: "string", array: true, nargs: 1, describe: "Variable substitution: --variable KEY=VALUE. Use ${KEY} in SQL." }) .option("set", { type: "string", array: true, nargs: 1, describe: "Query hint: --set KEY=VALUE (e.g. --set cz.sql.timezone=UTC)" }) .option("job-profile", { type: "string", describe: "Fetch execution profile for a completed job ID (separate from running SQL)" }) diff --git a/packages/cz-cli/src/main.ts b/packages/cz-cli/src/main.ts index 2e1871064..8a42be54f 100755 --- a/packages/cz-cli/src/main.ts +++ b/packages/cz-cli/src/main.ts @@ -1,5 +1,6 @@ #!/usr/bin/env bun // Dev entry point (bun run src/main.ts). The compiled binary uses src/bootstrap/boot.ts. +import { hasActiveSqlJobs } from "./commands/sql-lifecycle.js" import { checkAndUpdate } from "./auto-update.js" import { runCliWithTracking } from "./run-cli.js" import { createTraceparent } from "@clickzetta/sdk" @@ -10,8 +11,10 @@ if (!process.env.CLICKZETTA_TRACEPARENT) { } process.on("SIGINT", () => { + // Active SQL owns graceful cancellation before exiting. + if (hasActiveSqlJobs()) return const outputArgs = parseOutputArgs(process.argv.slice(2)) - // renderErrorOutput, matching sql.ts's SIGINT handler and error(): one error shape + // Match the active SQL handler and error(): one error shape // per format, so Ctrl-C under --format text is an ERROR row like any other failure. process.stdout.write(renderErrorOutput({ error: { code: "ABORTED", message: "Operation aborted by user." } }, outputArgs.format, outputArgs.field) + "\n") process.exit(130) diff --git a/packages/cz-cli/src/sql/cleanup-scope.ts b/packages/cz-cli/src/sql/cleanup-scope.ts new file mode 100644 index 000000000..858d19a9d --- /dev/null +++ b/packages/cz-cli/src/sql/cleanup-scope.ts @@ -0,0 +1,314 @@ +import { z } from "zod" +import { createServer, createConnection, type Socket } from "node:net" +import { + abortable, + abortAfter, + cancelJobAndWait, + staticTokenSource, + type ClientOptions, + type JobID, +} from "@clickzetta/sdk" + +export const CLEANUP_ENV = "CZ_SQL_CLEANUP" +const registration = z.object({ + type: z.literal("register"), + secret: z.string(), + job: z.object({ + id: z.string().min(1).max(256), + workspace: z.string().min(1), + instanceId: z.number().int().positive(), + }), + baseUrl: z + .string() + .url() + .refine((url) => ["http:", "https:"].includes(new URL(url).protocol)), + credential: z.object({ + token: z.string(), + instanceId: z.number(), + userId: z.number(), + headers: z.record(z.string(), z.string()).optional(), + }), + customHeaders: z.record(z.string(), z.string()).optional(), + timeoutMs: z.number().positive().finite(), +}) + +type Connection = { + socket: Socket + seen: number + expires: number + closing: boolean + entry?: { key: string; job: JobID; client: ClientOptions } +} + +/** Own jobs independently of shell hooks. One authenticated connection owns one job. */ +export async function createSqlSupervisor(options: { + heartbeatMs?: number + heartbeatTimeoutMs?: number + cleanupTimeoutMs?: number + onWarning: (warning: { jobId: string; reason: string }) => Promise +}) { + const secret = crypto.randomUUID() + const heartbeatMs = options.heartbeatMs ?? 1000 + const heartbeatTimeoutMs = options.heartbeatTimeoutMs ?? 10000 + const sockets = new Set() + const jobs = new Set() + const pending = new Set>() + let closing = false + let completion: Promise | undefined + + function cleanup(connection: Connection) { + if (connection.closing) { + connection.socket.destroy() + return + } + connection.closing = true + sockets.delete(connection) + connection.socket.destroy() + const entry = connection.entry + if (!entry) return + const work = (async () => { + try { + const result = await cancelJobAndWait(entry.client, entry.job, options.cleanupTimeoutMs ?? 5000) + if (!result.confirmed) await options.onWarning({ jobId: entry.job.id, reason: result.reason }) + } finally { + jobs.delete(entry.key) + } + })() + pending.add(work) + // Reporting must not create an unhandled rejection or stop other jobs' cleanup. + void work.catch(() => {}).finally(() => pending.delete(work)) + } + + const server = createServer((socket) => { + if (closing || sockets.size + pending.size >= 256) return socket.destroy() + const connection: Connection = { socket, seen: performance.now(), expires: Infinity, closing: false } + sockets.add(connection) + socket.on("error", () => cleanup(connection)) + socket.on("close", () => { + sockets.delete(connection) + cleanup(connection) + }) + readMessages(socket, (message) => { + if (closing || connection.closing) return cleanup(connection) + const parsed = z.object({ type: z.string() }).passthrough().safeParse(message) + if (!parsed.success) return cleanup(connection) + if (!connection.entry) { + const result = registration.safeParse(parsed.data) + if (!result.success || result.data.secret !== secret) return cleanup(connection) + const key = JSON.stringify([result.data.job.instanceId, result.data.job.workspace, result.data.job.id]) + if (jobs.has(key)) return cleanup(connection) + jobs.add(key) + connection.entry = { + key, + job: result.data.job, + client: { + baseUrl: result.data.baseUrl, + tokens: staticTokenSource(result.data.credential), + customHeaders: result.data.customHeaders, + }, + } + connection.expires = performance.now() + result.data.timeoutMs + connection.seen = performance.now() + socket.write(JSON.stringify({ type: "registered", heartbeatMs, heartbeatTimeoutMs }) + "\n") + return + } + // Expiry wins over a late heartbeat or handoff. + if (performance.now() >= connection.expires || performance.now() - connection.seen >= heartbeatTimeoutMs) + return cleanup(connection) + if (parsed.data.type === "heartbeat") { + connection.seen = performance.now() + socket.write('{"type":"heartbeat"}\n') + return + } + if (parsed.data.type === "release") { + jobs.delete(connection.entry.key) + connection.entry = undefined + connection.closing = true + socket.end('{"type":"released"}\n') + return + } + cleanup(connection) + }) + }) + await new Promise((resolve, reject) => { + server.once("error", reject) + server.listen(0, "127.0.0.1", () => { + server.removeListener("error", reject) + resolve() + }) + }) + const address = server.address() + if (!address || typeof address === "string") throw new Error("SQL supervisor did not bind a TCP port") + const timer = setInterval( + () => { + const now = performance.now() + for (const socket of sockets) { + if (now >= socket.expires || now - socket.seen >= heartbeatTimeoutMs) cleanup(socket) + } + }, + Math.min(heartbeatMs, 1000), + ) + timer.unref() + server.unref() + return { + env: { [CLEANUP_ENV]: JSON.stringify({ url: `tcp://127.0.0.1:${address.port}`, secret }) }, + close() { + if (completion) return completion + closing = true + clearInterval(timer) + for (const socket of sockets) cleanup(socket) + completion = Promise.allSettled([...pending]).then( + () => + new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + }), + ) + return completion + }, + } +} + +function parseMessage(data: string): unknown { + try { + return JSON.parse(data) + } catch { + return undefined + } +} + +/** Registration is acknowledged before submit; disconnect aborts local execution. */ +export async function registerSqlCleanup( + client: ClientOptions, + job: JobID, + signal: AbortSignal, + timeoutMs: number, + onLost: () => void, +) { + const raw = process.env[CLEANUP_ENV] + if (!raw) return + if (raw === "unavailable") + throw new Error( + "SQL cleanup supervisor unavailable; query was not submitted. Check local socket and ~/.clickzetta write permissions.", + ) + const address = z.object({ url: z.string().url(), secret: z.string().uuid() }).parse(JSON.parse(raw)) + const url = new URL(address.url) + if (url.hostname !== "127.0.0.1" || url.protocol !== "tcp:" || !url.port) + throw new Error("Invalid SQL cleanup supervisor") + const deadline = abortAfter(2000, signal) + try { + const credential = await abortable(client.tokens.get(), deadline.signal) + deadline.signal.throwIfAborted() + const socket = createConnection({ host: url.hostname, port: Number(url.port) }) + const ready = Promise.withResolvers() + const released = Promise.withResolvers() + // Handoff may never be requested. Still observe rejection if the socket dies. + void released.promise.catch(() => {}) + let acknowledged = false + let done = false + let timer: ReturnType | undefined + let lastSeen = performance.now() + const abandon = () => { + if (timer) clearInterval(timer) + socket.destroy() + } + const lost = () => { + if (done) return + done = true + abandon() + const error = new Error(`SQL job ${job.id}: cleanup supervisor connection lost`) + ready.reject(error) + released.reject(error) + if (acknowledged) onLost() + } + socket.on("connect", () => { + socket.write( + JSON.stringify({ + type: "register", + secret: address.secret, + job, + baseUrl: client.baseUrl, + credential, + customHeaders: client.customHeaders, + timeoutMs, + }) + "\n", + ) + }) + socket.on("error", lost) + socket.on("close", lost) + readMessages(socket, (message) => { + const parsed = z + .object({ + type: z.enum(["registered", "heartbeat", "released"]), + heartbeatMs: z.number().positive().optional(), + heartbeatTimeoutMs: z.number().positive().optional(), + }) + .safeParse(message) + if (!parsed.success || done) return lost() + if (parsed.data.type === "registered" && !acknowledged) { + const heartbeatMs = parsed.data.heartbeatMs + const heartbeatTimeoutMs = parsed.data.heartbeatTimeoutMs + if (!heartbeatMs || !heartbeatTimeoutMs) return lost() + acknowledged = true + timer = setInterval(() => { + if (performance.now() - lastSeen >= heartbeatTimeoutMs) return lost() + if (!socket.destroyed) socket.write('{"type":"heartbeat"}\n') + }, heartbeatMs) + ready.resolve() + return + } + if (parsed.data.type === "heartbeat") { + lastSeen = performance.now() + return + } + if (parsed.data.type === "released") { + done = true + if (timer) clearInterval(timer) + released.resolve() + socket.destroy() + return + } + lost() + }) + try { + await abortable(ready.promise, deadline.signal) + } catch (error) { + done = true + abandon() + throw error + } + return { + abandon, + async release() { + if (done || socket.destroyed) throw new Error("SQL cleanup acknowledgement failed") + const limit = abortAfter(1000) + try { + socket.write('{"type":"release"}\n') + await abortable(released.promise, limit.signal) + } finally { + limit.dispose() + abandon() + } + }, + } + } finally { + deadline.dispose() + } +} + +// Bound each newline-delimited UTF-8 frame before parsing; preserve split multibyte characters. +function readMessages(socket: Socket, receive: (message: unknown) => void) { + let buffer = Buffer.alloc(0) + socket.on("data", (chunk: Buffer) => { + buffer = Buffer.concat([buffer, chunk]) + while (!socket.destroyed) { + if (socket.writableLength > 64 * 1024) return socket.destroy() + const end = buffer.indexOf(10) + if (end < 0) break + if (end > 64 * 1024) return socket.destroy() + const message = parseMessage(buffer.subarray(0, end).toString("utf8")) + buffer = buffer.subarray(end + 1) + receive(message) + } + if (buffer.length > 64 * 1024) socket.destroy() + }) +} diff --git a/packages/cz-cli/src/sql/supervisor-runtime.ts b/packages/cz-cli/src/sql/supervisor-runtime.ts new file mode 100644 index 000000000..5f06533c0 --- /dev/null +++ b/packages/cz-cli/src/sql/supervisor-runtime.ts @@ -0,0 +1,32 @@ +import { appendFile, mkdir } from "node:fs/promises" +import { homedir } from "node:os" +import path from "node:path" +import { CLEANUP_ENV, createSqlSupervisor } from "./cleanup-scope.js" + +/** Bootstrap before importing the agent or creating its Worker: both inherit the endpoint. */ +export async function withSqlSupervisor(run: () => Promise): Promise { + const previous = process.env[CLEANUP_ENV] + const directory = path.join(process.env.CLICKZETTA_TEST_HOME || homedir(), ".clickzetta") + const logfile = path.join(directory, "sql-cleanup.jsonl") + const supervisor = await (async () => { + // Open diagnostics before admitting jobs, not in a failing finalizer. + await mkdir(directory, { recursive: true }) + await appendFile(logfile, "", { mode: 0o600 }) + return createSqlSupervisor({ + async onWarning(warning) { + // Never write into a terminal owned by the TUI renderer or persist credentials. + await appendFile(logfile, JSON.stringify({ time: new Date().toISOString(), ...warning }) + "\n") + }, + }) + })().catch(() => undefined) + // Keep non-SQL tools usable if local sockets/diagnostics are unavailable. + // SQL must still fail closed rather than silently run without supervision. + process.env[CLEANUP_ENV] = supervisor?.env[CLEANUP_ENV] ?? "unavailable" + try { + return await run() + } finally { + if (previous === undefined) delete process.env[CLEANUP_ENV] + else process.env[CLEANUP_ENV] = previous + await supervisor?.close() + } +} diff --git a/packages/cz-cli/test/sql-cancellation.test.ts b/packages/cz-cli/test/sql-cancellation.test.ts new file mode 100644 index 000000000..071db99ed --- /dev/null +++ b/packages/cz-cli/test/sql-cancellation.test.ts @@ -0,0 +1,171 @@ +import { expect, test } from "bun:test" + +test("query deadline aborts a blocked submit and confirms cancellation", async () => { + const cancelled: string[] = [] + const submitted: string[] = [] + const server = Bun.serve({ + port: 0, + async fetch(request) { + const body = await request.json() + if (request.url.endsWith("submitJob")) { + submitted.push(body.jobDesc.jobId.id) + expect(body.jobDesc.jobTimeoutMs).toBe(150) + return new Promise(() => {}) + } + if (request.url.endsWith("cancelJob")) { + cancelled.push(body.job_id.id) + return Response.json({}) + } + return Response.json({ status: { state: "CANCELLED" } }) + }, + }) + const child = Bun.spawn( + [ + process.execPath, + "--eval", + `import { execSql } from ${JSON.stringify(new URL("../src/commands/exec.ts", import.meta.url).pathname)}; + import { anonymous } from '@clickzetta/sdk'; + await execSql({ config: { workspace: 'ws', schema: 'public', vcluster: 'vc', instance: 'inst' }, + clientOpts: { baseUrl: ${JSON.stringify(server.url.origin)}, tokens: anonymous() }, instanceId: () => 1 }, + 'select 1', { timeoutMs: 150 });`, + ], + { stdout: "pipe", stderr: "pipe" }, + ) + try { + expect(await child.exited).not.toBe(0) + expect(submitted).toHaveLength(1) + expect(cancelled).toEqual(submitted) + expect(await new Response(child.stderr).text()).toContain("interrupted or timed out") + } finally { + child.kill("SIGKILL") + await child.exited + await server.stop(true) + } +}) + +// Real child processes and HTTP requests: signal handling must work outside the +// test process, including while a HYBRID submission has not returned yet. +for (const signal of ["SIGTERM", "SIGINT"] as const) { + for (const phase of ["submit", "poll"] as const) { + test(`${signal} cancels all active jobs during ${phase}`, async () => { + const submitted: string[] = [] + const cancelled: string[] = [] + const ready = Promise.withResolvers() + let polls = 0 + const server = Bun.serve({ + port: 0, + async fetch(request) { + const body = await request.json() + if (request.url.endsWith("/lh/submitJob")) { + submitted.push(body.jobDesc.jobId.id) + expect(body.jobDesc.jobTimeoutMs).toBe(30000) + if (phase === "submit") { + if (submitted.length === 2) ready.resolve() + return new Promise(() => {}) + } + return Response.json({ status: { state: "RUNNING" } }) + } + if (request.url.endsWith("/lh/cancelJob")) { + cancelled.push(body.job_id.id) + return Response.json({ code: "SUCCESS", data: {} }) + } + if (++polls === 2) ready.resolve() + return Response.json({ status: { state: cancelled.length ? "CANCELLED" : "RUNNING" } }) + }, + }) + const child = Bun.spawn( + [ + process.execPath, + "--eval", + ` + import { execSql } from ${JSON.stringify(new URL("../src/commands/exec.ts", import.meta.url).pathname)}; + import { anonymous } from "@clickzetta/sdk"; + const ctx = { + config: { workspace: "ws", schema: "public", vcluster: "vc", instance: "inst" }, + clientOpts: { baseUrl: ${JSON.stringify(server.url.href)}, tokens: anonymous() }, + instanceId: () => 1, + }; + await Promise.all([execSql(ctx, "select 1", { timeoutMs: 30000 }), execSql(ctx, "select 2", { timeoutMs: 30000 })]); + `, + ], + { stdout: "pipe", stderr: "pipe" }, + ) + try { + await Promise.race([ + ready.promise, + Bun.sleep(5000).then(() => { + throw new Error("SQL did not start") + }), + ]) + child.kill(signal) + expect(await child.exited).toBe(signal === "SIGTERM" ? 143 : 130) + expect(cancelled.sort()).toEqual(submitted.sort()) + expect(cancelled).toHaveLength(2) + expect(await new Response(child.stdout).text()).toContain("ABORTED") + } finally { + child.kill("SIGKILL") + server.stop(true) + } + }, 10000) + } +} + +test("detached async jobs release signal handlers", async () => { + const child = Bun.spawn( + [ + process.execPath, + "--eval", + ` + import { execSql } from ${JSON.stringify(new URL("../src/commands/exec.ts", import.meta.url).pathname)}; + import { hasActiveSqlJobs } from ${JSON.stringify(new URL("../src/commands/sql-lifecycle.ts", import.meta.url).pathname)}; + import { anonymous } from "@clickzetta/sdk"; + const server = Bun.serve({ port: 0, fetch: () => Response.json({ status: { state: "RUNNING" } }) }); + const before = process.listenerCount("SIGTERM"); + await execSql({ config: { workspace: "ws", schema: "public", vcluster: "vc", instance: "inst" }, clientOpts: { baseUrl: server.url.href, tokens: anonymous() }, instanceId: () => 1 }, "select 1", { asynchronous: true }); + console.log(JSON.stringify({ active: hasActiveSqlJobs(), listeners: process.listenerCount("SIGTERM") - before })); + server.stop(true); + `, + ], + { stdout: "pipe", stderr: "pipe" }, + ) + expect(await child.exited).toBe(0) + expect(JSON.parse(await new Response(child.stdout).text())).toEqual({ active: false, listeners: 0 }) +}) + +test("a stalled cancellation cannot block process shutdown", async () => { + const ready = Promise.withResolvers() + const server = Bun.serve({ + port: 0, + fetch(request) { + if (request.url.endsWith("/lh/cancelJob")) return new Promise(() => {}) + ready.resolve() + return Response.json({ status: { state: "RUNNING" } }) + }, + }) + const child = Bun.spawn( + [ + process.execPath, + "--eval", + ` + import { execSql } from ${JSON.stringify(new URL("../src/commands/exec.ts", import.meta.url).pathname)}; + import { anonymous } from "@clickzetta/sdk"; + await execSql({ config: { workspace: "ws", schema: "public", vcluster: "vc", instance: "inst" }, clientOpts: { baseUrl: ${JSON.stringify(server.url.href)}, tokens: anonymous() }, instanceId: () => 1 }, "select 1"); + `, + ], + { stdout: "pipe", stderr: "pipe" }, + ) + try { + await Promise.race([ + ready.promise, + Bun.sleep(5000).then(() => { + throw new Error("SQL did not start") + }), + ]) + child.kill("SIGTERM") + expect(await child.exited).toBe(143) + expect(await new Response(child.stderr).text()).toContain("cancellation unconfirmed") + } finally { + child.kill("SIGKILL") + server.stop(true) + } +}, 10000) diff --git a/packages/cz-cli/test/sql-cleanup-scope.test.ts b/packages/cz-cli/test/sql-cleanup-scope.test.ts new file mode 100644 index 000000000..b909c8a4f --- /dev/null +++ b/packages/cz-cli/test/sql-cleanup-scope.test.ts @@ -0,0 +1,443 @@ +import { Effect } from "effect" +import { ChildProcess } from "effect/unstable/process" +import { AppProcess } from "@opencode-ai/core/process" +import { createConnection, createServer } from "node:net" +import { expect, test } from "bun:test" +import { mkdtemp, rm } from "node:fs/promises" +import os from "node:os" +import path from "node:path" +import { pathToFileURL } from "node:url" +import { createSqlSupervisor } from "../src/sql/cleanup-scope.js" + +const execModule = new URL("../src/commands/exec.ts", import.meta.url).pathname +function program(url: string, options = "{ timeoutMs: 30000 }") { + return `import { execSql } from ${JSON.stringify(execModule)}; + import { anonymous } from '@clickzetta/sdk'; + const ctx = { config: { workspace: 'ws', schema: 'public', vcluster: 'vc', instance: 'inst' }, + clientOpts: { baseUrl: ${JSON.stringify(url)}, tokens: anonymous() }, instanceId: () => 1 }; + await execSql(ctx, 'select 1', ${options});` +} + +async function until(predicate: () => boolean, timeout = 5000) { + const deadline = performance.now() + timeout + while (!predicate()) { + if (performance.now() >= deadline) throw new Error("condition did not become true") + await Bun.sleep(10) + } +} + +function queryServer(stage: "submit" | "poll" = "poll") { + const submitted: string[] = [] + const cancelled: string[] = [] + const active = new Set() + const timeouts: (number | undefined)[] = [] + const server = Bun.serve({ + port: 0, + async fetch(request) { + const body = await request.json() + if (request.url.endsWith("submitJob")) { + const id = body.jobDesc.jobId.id + submitted.push(id) + timeouts.push(body.jobDesc.jobTimeoutMs) + active.add(id) + if (stage === "submit") return new Promise(() => {}) + return Response.json({ status: { state: "RUNNING" } }) + } + if (request.url.endsWith("cancelJob")) { + cancelled.push(body.job_id.id) + active.delete(body.job_id.id) + return Response.json({}) + } + return Response.json({ status: { state: active.has(body.getResultRequest.jobId.id) ? "RUNNING" : "CANCELLED" } }) + }, + }) + return { server, submitted, cancelled, active, timeouts } +} + +for (const stage of ["submit", "poll"] as const) { + test(`SIGKILL during ${stage} triggers disconnect cleanup without touching another connection`, async () => { + const remote = queryServer(stage) + const warnings: unknown[] = [] + const supervisor = await createSqlSupervisor({ + onWarning: async (warning) => { + warnings.push(warning) + }, + }) + const children = [1, 2].map(() => + Bun.spawn([process.execPath, "--eval", program(remote.server.url.origin)], { + env: { ...process.env, ...supervisor.env }, + stdout: "pipe", + stderr: "pipe", + }), + ) + try { + await until(() => remote.submitted.length === 2) + children[0].kill("SIGKILL") + await children[0].exited + // No shell hook, no explicit supervisor.close(): socket death drives cleanup. + await until(() => remote.cancelled.length === 1) + expect(remote.active.size).toBe(1) + await Bun.sleep(100) + expect(remote.cancelled).toHaveLength(1) + children[1].kill("SIGKILL") + await children[1].exited + await until(() => remote.active.size === 0) + expect(new Set(remote.cancelled).size).toBe(2) + expect(warnings).toEqual([]) + } finally { + children.forEach((child) => child.kill("SIGKILL")) + await Promise.all(children.map((child) => child.exited)) + await supervisor.close() + await remote.server.stop(true) + } + }, 15000) +} + +test("explicit async handoff survives child exit and supervisor shutdown", async () => { + const remote = queryServer() + const supervisor = await createSqlSupervisor({ onWarning: async () => {} }) + const child = Bun.spawn( + [process.execPath, "--eval", program(remote.server.url.origin, "{ timeoutMs: 30000, asynchronous: true }")], + { + env: { ...process.env, ...supervisor.env }, + stdout: "pipe", + stderr: "pipe", + }, + ) + try { + const code = await child.exited + expect({ code, error: code ? await new Response(child.stderr).text() : "" }).toEqual({ code: 0, error: "" }) + await supervisor.close() + expect(remote.cancelled).toEqual([]) + expect(remote.timeouts).toEqual([30000]) + } finally { + child.kill("SIGKILL") + await child.exited + await supervisor.close() + await remote.server.stop(true) + } +}) + +test("unavailable supervisor prevents submission", async () => { + const remote = queryServer() + const supervisor = await createSqlSupervisor({ onWarning: async () => {} }) + await supervisor.close() + const child = Bun.spawn([process.execPath, "--eval", program(remote.server.url.origin)], { + env: { ...process.env, ...supervisor.env }, + stdout: "pipe", + stderr: "pipe", + }) + try { + expect(await child.exited).not.toBe(0) + expect(remote.submitted).toEqual([]) + } finally { + child.kill("SIGKILL") + await child.exited + await remote.server.stop(true) + } +}) + +async function register(env: { CZ_SQL_CLEANUP: string }, baseUrl: string, id = "owned", timeoutMs = 30000) { + const address = JSON.parse(env.CZ_SQL_CLEANUP) + const url = new URL(address.url) + const socket = createConnection({ host: url.hostname, port: Number(url.port) }) + const ready = Promise.withResolvers() + const timer = setTimeout(() => { + socket.destroy() + ready.reject(new Error("registration timed out")) + }, 2000) + socket.on("connect", () => + socket.write( + JSON.stringify({ + type: "register", + secret: address.secret, + job: { id, workspace: "ws", instanceId: 1 }, + baseUrl, + credential: { token: "private-token", instanceId: 1, userId: 1 }, + timeoutMs, + }) + "\n", + ), + ) + socket.on("data", () => ready.resolve()) + socket.on("close", () => ready.reject(new Error("registration rejected"))) + socket.on("error", () => ready.reject(new Error("registration failed"))) + try { + await ready.promise + return socket + } finally { + clearTimeout(timer) + } +} + +for (const expiry of ["heartbeat", "deadline"] as const) { + test(`${expiry} expiry cancels a still-connected job`, async () => { + const remote = queryServer() + remote.active.add("owned") + const supervisor = await createSqlSupervisor({ + heartbeatMs: 20, + heartbeatTimeoutMs: expiry === "heartbeat" ? 100 : 1000, + onWarning: async () => {}, + }) + const socket = await register( + supervisor.env, + remote.server.url.origin, + "owned", + expiry === "deadline" ? 100 : 30000, + ) + try { + await until(() => remote.cancelled.length > 0) + expect(remote.cancelled).toEqual(["owned"]) + expect(remote.active.size).toBe(0) + } finally { + socket.destroy() + await supervisor.close() + await remote.server.stop(true) + } + }) +} + +test("healthy heartbeats keep a long-running query alive", async () => { + const remote = queryServer("submit") + const supervisor = await createSqlSupervisor({ heartbeatMs: 20, heartbeatTimeoutMs: 300, onWarning: async () => {} }) + const child = Bun.spawn([process.execPath, "--eval", program(remote.server.url.origin)], { + env: { ...process.env, ...supervisor.env }, + stdout: "pipe", + stderr: "pipe", + }) + try { + await until(() => remote.submitted.length === 1) + await Bun.sleep(700) + expect(remote.cancelled).toEqual([]) + child.kill("SIGKILL") + await child.exited + await until(() => remote.cancelled.length === 1) + } finally { + child.kill("SIGKILL") + await child.exited + await supervisor.close() + await remote.server.stop(true) + } +}) + +test("wrong capability cannot register a job", async () => { + const remote = queryServer() + const supervisor = await createSqlSupervisor({ onWarning: async () => {} }) + const address = JSON.parse(supervisor.env.CZ_SQL_CLEANUP) + try { + await expect( + register( + { CZ_SQL_CLEANUP: JSON.stringify({ ...address, secret: crypto.randomUUID() }) }, + remote.server.url.origin, + ), + ).rejects.toThrow("rejected") + expect(remote.cancelled).toEqual([]) + } finally { + await supervisor.close() + await remote.server.stop(true) + } +}) + +test("cleanup reports unconfirmed job IDs within its budget without leaking credentials", async () => { + const warnings: { jobId: string; reason: string }[] = [] + const server = Bun.serve({ port: 0, fetch: () => new Promise(() => {}) }) + const supervisor = await createSqlSupervisor({ + cleanupTimeoutMs: 150, + onWarning: async (warning) => { + warnings.push(warning) + }, + }) + const socket = await register(supervisor.env, server.url.origin) + try { + const start = performance.now() + socket.destroy() + await until(() => warnings.length === 1) + expect(performance.now() - start).toBeLessThan(1500) + expect(warnings[0].jobId).toBe("owned") + expect(JSON.stringify(warnings)).not.toContain("private-token") + await supervisor.close() + await supervisor.close() + } finally { + socket.destroy() + await supervisor.close() + await server.stop(true) + } +}) + +test("lost async acknowledgement fails the command and cancels the submitted job", async () => { + const remote = queryServer() + const server = createServer((socket) => { + socket.on("error", () => {}) + socket.on("data", (data) => { + for (const line of data.toString().trim().split("\n")) { + const message = JSON.parse(line) + if (message.type === "register") + socket.write(JSON.stringify({ type: "registered", heartbeatMs: 1000, heartbeatTimeoutMs: 10000 }) + "\n") + if (message.type === "release") socket.destroy() + } + }) + }) + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)) + const address = server.address() + if (!address || typeof address === "string") throw new Error("missing address") + const child = Bun.spawn( + [process.execPath, "--eval", program(remote.server.url.origin, "{ timeoutMs: 30000, asynchronous: true }")], + { + env: { + ...process.env, + CZ_SQL_CLEANUP: JSON.stringify({ url: `tcp://127.0.0.1:${address.port}`, secret: crypto.randomUUID() }), + }, + stdout: "pipe", + stderr: "pipe", + }, + ) + try { + expect(await child.exited).not.toBe(0) + expect(remote.submitted).toHaveLength(1) + expect(remote.cancelled).toEqual(remote.submitted) + const stderr = await new Response(child.stderr).text() + expect(stderr).toContain("async handoff acknowledgement failed") + expect(stderr).toContain(remote.submitted[0]) + } finally { + child.kill("SIGKILL") + await child.exited + await new Promise((resolve) => server.close(() => resolve())) + await remote.server.stop(true) + } +}) + +test("bundled cz bootstrap passes supervision to children through inherited environment", async () => { + const directory = await mkdtemp(path.join(os.tmpdir(), "cz-supervisor-")) + const remote = queryServer("submit") + const previous = process.env.CZ_SQL_CLEANUP + try { + const build = await Bun.build({ + entrypoints: [new URL("../src/sql/supervisor-runtime.ts", import.meta.url).pathname], + target: "bun", + format: "esm", + outdir: directory, + }) + expect(build.success).toBe(true) + const runtime: typeof import("../src/sql/supervisor-runtime.js") = await import( + pathToFileURL(build.outputs[0].path).href + ) + let endpoint = "" + await runtime.withSqlSupervisor(async () => { + endpoint = JSON.parse(process.env.CZ_SQL_CLEANUP!).url + const child = Bun.spawn([process.execPath, "--eval", program(remote.server.url.origin)], { + env: { ...process.env }, + stdout: "pipe", + stderr: "pipe", + }) + try { + await until(() => remote.submitted.length === 1) + child.kill("SIGKILL") + await child.exited + await until(() => remote.cancelled.length === 1) + } finally { + child.kill("SIGKILL") + await child.exited + } + }) + expect(process.env.CZ_SQL_CLEANUP).toBe(previous) + const url = new URL(endpoint) + const probe = createConnection({ host: url.hostname, port: Number(url.port) }) + await new Promise((resolve) => probe.once("error", () => resolve())) + probe.destroy() + } finally { + await remote.server.stop(true) + await rm(directory, { recursive: true, force: true }) + } +}, 15000) + +test("unmodified upstream process runner inherits supervision and cancels after timeout", async () => { + const { withSqlSupervisor } = await import("../src/sql/supervisor-runtime.js") + const remote = queryServer("submit") + try { + await withSqlSupervisor(async () => { + await Effect.runPromise( + Effect.gen(function* () { + const processService = yield* AppProcess.Service + yield* processService + .run( + ChildProcess.make(process.execPath, ["--eval", program(remote.server.url.origin)], { + killSignal: "SIGKILL", + }), + { timeout: "1 second" }, + ) + .pipe(Effect.ignore) + }).pipe(Effect.provide(AppProcess.defaultLayer)), + ) + expect(remote.submitted).toHaveLength(1) + await until(() => remote.cancelled.length === 1) + }) + } finally { + await remote.server.stop(true) + } +}) + +test("existing cz Worker environment bridge carries the supervisor endpoint", async () => { + const directory = await mkdtemp(path.join(os.tmpdir(), "cz-supervisor-worker-")) + await Bun.write(path.join(directory, "worker.ts"), "postMessage(process.env.CZ_SQL_CLEANUP); close();") + const child = Bun.spawn( + [ + process.execPath, + "--eval", + ` + import { withSqlSupervisor } from ${JSON.stringify(new URL("../src/sql/supervisor-runtime.ts", import.meta.url).pathname)}; + import { installClickzettaWorkerEnvShim } from ${JSON.stringify(new URL("../src/bootstrap/runtime-config.ts", import.meta.url).pathname)}; + await withSqlSupervisor(async () => { + installClickzettaWorkerEnvShim(); + const worker = new Worker(${JSON.stringify(path.join(directory, "worker.ts"))}); + try { + const value = await new Promise((resolve, reject) => { worker.onmessage = e => resolve(e.data); worker.onerror = reject; }); + console.log(JSON.stringify({ inherited: value === process.env.CZ_SQL_CLEANUP, protocol: JSON.parse(value).url.startsWith('tcp:') })); + } finally { await worker.terminate(); } + }); + `, + ], + { env: { ...process.env }, stdout: "pipe", stderr: "pipe" }, + ) + try { + const code = await child.exited + expect({ code, error: code ? await new Response(child.stderr).text() : "" }).toEqual({ code: 0, error: "" }) + expect(JSON.parse(await new Response(child.stdout).text())).toEqual({ inherited: true, protocol: true }) + } finally { + child.kill("SIGKILL") + await child.exited + await rm(directory, { recursive: true, force: true }) + } +}) + +test("supervisor startup failure leaves non-SQL work usable but blocks SQL admission", async () => { + const directory = await mkdtemp(path.join(os.tmpdir(), "cz-supervisor-unavailable-")) + await Bun.write(path.join(directory, ".clickzetta"), "not a directory") + const remote = queryServer() + const child = Bun.spawn( + [ + process.execPath, + "--eval", + ` + import { withSqlSupervisor } from ${JSON.stringify(new URL("../src/sql/supervisor-runtime.ts", import.meta.url).pathname)}; + await withSqlSupervisor(async () => { + console.log('non-sql-work-ran'); + const { execSql } = await import(${JSON.stringify(execModule)}); + const { anonymous } = await import('@clickzetta/sdk'); + await execSql({ config: { workspace: 'ws', schema: 'public', vcluster: 'vc', instance: 'inst' }, + clientOpts: { baseUrl: ${JSON.stringify(remote.server.url.origin)}, tokens: anonymous() }, instanceId: () => 1 }, 'select 1'); + }); + `, + ], + { env: { ...process.env, CLICKZETTA_TEST_HOME: directory }, stdout: "pipe", stderr: "pipe" }, + ) + try { + expect(await child.exited).not.toBe(0) + expect(await new Response(child.stdout).text()).toContain("non-sql-work-ran") + expect(await new Response(child.stderr).text()).toContain("SQL cleanup supervisor unavailable") + expect(remote.submitted).toEqual([]) + } finally { + child.kill("SIGKILL") + await child.exited + await remote.server.stop(true) + await rm(directory, { recursive: true, force: true }) + } +}) diff --git a/packages/cz-cli/test/sql-timeout-options.test.ts b/packages/cz-cli/test/sql-timeout-options.test.ts new file mode 100644 index 000000000..f0978caa3 --- /dev/null +++ b/packages/cz-cli/test/sql-timeout-options.test.ts @@ -0,0 +1,45 @@ +import { beforeEach, expect, test } from "bun:test" +import { join } from "node:path" +import { onFetch, requireTestHome, sqlSuccess, stubStudioContext } from "./support/cz-fixtures.js" + +const { execute } = await import("../src/execute.ts") +const timeouts: (number | undefined)[] = [] +beforeEach(async () => { + timeouts.length = 0 + stubStudioContext() + await Bun.file(join(requireTestHome(), ".clickzetta", "profiles.toml")).write( + 'default_profile = "test"\n[profiles.test]\npat = "pat"\nworkspace = "ws0"\ninstance = "inst"\n', + ) + onFetch({ + match: (url) => url.includes("/lh/submitJob"), + respond: (_url, _method, body) => { + timeouts.push((body as { jobDesc: { jobTimeoutMs?: number } }).jobDesc.jobTimeoutMs) + return sqlSuccess(["v"], [[1]]) + }, + }) +}) + +for (const mode of ["--sync", "--async"]) { + for (const timeout of ["0", "-1", "NaN", "Infinity"]) { + test(`sql ${mode} rejects --timeout ${timeout} before submission`, async () => { + const result = await execute("sql", ["select 1", mode, "--timeout", timeout]) + expect(result.exitCode).not.toBe(0) + expect(result.output).toContain("--timeout") + expect(timeouts).toEqual([]) + }) + } + test(`sql ${mode} propagates explicit timeout to the server`, async () => { + const result = await execute("sql", ["select 1", mode, "--timeout", "12"]) + expect(result.exitCode).toBe(0) + expect(timeouts).toEqual([12000]) + }) +} + +for (const command of ["table list", "schema list"]) { + test(`${command} does not acquire a new default five-minute timeout`, async () => { + const result = await execute(command) + expect(result.exitCode).toBe(0) + expect(timeouts.length).toBeGreaterThan(0) + expect(timeouts.every((value) => value === undefined)).toBe(true) + }) +} diff --git a/specs/sql-cancellation.md b/specs/sql-cancellation.md new file mode 100644 index 000000000..9d7be345e --- /dev/null +++ b/specs/sql-cancellation.md @@ -0,0 +1,97 @@ +# SQL cancellation ownership + +Killing a shell process does not cancel the SQL it submitted. `execSql` owns a +remote job from before submission until terminal confirmation or an explicit +asynchronous handoff. This includes a HYBRID submission that has not returned. + +## No upstream patches + +All supervision lives in `packages/cz-cli`; transport cancellation lives in the +ClickZetta SDK. There are no changes to opencode, Core, TUI, their tools or their +plugin APIs. No command rewriting, stdout parsing or global shell interception +is used. + +The cz agent bootstrap starts a loopback TCP supervisor before importing the +agent runtime. Its endpoint and random capability are passed in `CZ_SQL_CLEANUP`. +The existing cz Worker environment bridge copies them into the TUI server Worker; +upstream shell tools inherit the current environment through their existing +process runner. A separate connection owns each SQL job, so concurrent commands +and sessions do not share job ownership. Remote servers use the supervisor on the +machine actually executing their CLI subprocesses, not the attached client. + +If the supervisor or its diagnostic destination cannot start, non-SQL tools can +still run. The environment explicitly marks supervision unavailable and SQL +admission fails closed. A missing inherited environment (for example after an +explicit `env -i`) uses standalone CLI behavior; it cannot promise supervisor +recovery. + +## Protocol and recovery + +Before submitting anything, a SQL process connects, sends its capability, job ID, +credential snapshot, endpoint and timeout, and waits for acknowledgement. Each +connection accepts one job. Registration failure prevents submission. Frames are +newline-delimited JSON, limited to 64 KiB, with at most 256 connections/cleanup +operations. Credentials remain in memory and never enter the command line, disk +or diagnostic output. Registration and control operations are scoped to the +connection; one connection cannot release another job. + +The child sends a heartbeat every second and expects replies. Disconnect triggers +supervisor cancellation immediately; a 10-second heartbeat lease handles frozen +processes or a connection that does not close promptly. The child also aborts +execution if it loses its supervisor. An independent supervisor deadline bounds +a job even while heartbeats continue. Both ends use monotonic elapsed time. + +Cancellation validates business status and polls for a terminal state. HTTP +success alone is insufficient. Missing/unsubmitted jobs are not terminal proof: +cleanup repeats cancellation to cover delayed submission visibility. Execution +cancellation stops HTTP, credential waits and submit/poll retries; cleanup uses a +fresh signal with an independent total budget. + +A child unregisters after confirmed completion/cancellation. `--async` explicitly +hands off ownership only after acknowledgement. Losing that acknowledgement +fails the command and attempts cancellation rather than reporting a successful +handoff. Disconnect after a successful handoff does not cancel the detached job. + +| Boundary | Budget | +| ----------------------- | --------------------------------------------- | +| Registration | 2 seconds, also subject to the query deadline | +| Child cleanup | 1.5 seconds per job, concurrently | +| Signal shutdown | Hard 2-second exit bound | +| Supervisor cleanup | 5 seconds per job, concurrently | +| Heartbeat expiry | 10 seconds, checked every second | +| Handoff acknowledgement | 1 second | + +Supervisor cleanup is asynchronous relative to the upstream shell result. A shell +may return before cancellation is confirmed; this design deliberately makes no +claim that shell return proves remote queue release. Unconfirmed cleanup records +job ID and sanitized reason in `~/.clickzetta/sql-cleanup.jsonl`, without writing +into a terminal owned by the TUI renderer. Normal bootstrap return closes the +supervisor and attempts cancellation of all outstanding registrations. + +## Timeout compatibility and limits + +Standalone `cz-cli sql` retains its default 300-second timeout and now actively +cancels on timeout, SIGINT or SIGTERM. Positive explicit timeouts are sent to the +server, including for `--async`. Invalid/non-positive CLI timeout values fail at +the command boundary; zero is not an unlimited-timeout sentinel. + +Standalone `execSql` callers without a timeout (such as table/schema/file commands) +keep their previous deployment-defined timeout. They still clean up on caught +failure and signals. Agent-supervised calls without an explicit timeout receive +a 300-second fallback. No upstream hook exposes the shell's remaining budget, so +supervision uses the SQL deadline plus disconnect/heartbeat detection instead. + +If the entire agent is killed, loses network access, or the credential snapshot +expires, immediate cancellation is not guaranteed. A finite server timeout is +the final fallback for supervised jobs. Ordinary coordinator jobs check timeout +from submission across queued and running states; deployed behavior still needs +live validation. Continuous jobs have different semantics. Cancellation cannot +undo completed writes and is not transaction rollback. No durable recovery is +introduced. + +Tests use real sockets, HTTP servers and subprocesses for SIGINT/SIGTERM/SIGKILL, +submit/poll races, heartbeat expiry, healthy heartbeats, connection isolation, +failed admission/handoff, bounded cleanup, bootstrap bundling, Worker environment +inheritance and the unchanged upstream process runner. CLI tests cover timeout +validation and defaults. Windows signal behavior and live deployment recovery +require environment-specific validation. From 3dc09d605ed74fdbace87f6771a36b506b4b33a5 Mon Sep 17 00:00:00 2001 From: suibianwanwan Date: Thu, 8 Oct 2026 11:46:29 +0800 Subject: [PATCH 4/5] fix(sql): address cancellation review findings - accept instanceId 0 so profile add --verify and setup work under supervision - keep supervision when the diagnostics log is unwritable; tolerate close errors - resolve credentials under the query deadline, not the 2s handshake budget - name abort causes (timeout, signal, lost supervisor) and keep the cause - keep job_id and the user-interrupt message in the signal envelope - send no server timeout for --async unless --timeout is explicit - accept proto3 cancel responses without an error status - validate --timeout after --job-profile, which ignores it - drop the outbound-buffer read check; compare secrets in constant time Co-Authored-By: Claude Opus 5.5 (1M context) --- packages/clickzetta-sdk/src/sql/cancel.ts | 6 +-- .../test/fixtures/cancellation.ts | 12 ++++++ packages/cz-cli/src/commands/exec.ts | 31 +++++++++----- packages/cz-cli/src/commands/sql-lifecycle.ts | 21 +++++++--- packages/cz-cli/src/commands/sql.ts | 20 +++++---- packages/cz-cli/src/sql/cleanup-scope.ts | 21 +++++++--- packages/cz-cli/src/sql/supervisor-runtime.ts | 32 +++++++------- packages/cz-cli/test/sql-cancellation.test.ts | 11 ++++- .../cz-cli/test/sql-cleanup-scope.test.ts | 42 +++++++++++++++---- .../cz-cli/test/sql-timeout-options.test.ts | 12 ++++++ specs/sql-cancellation.md | 24 +++++++---- 11 files changed, 168 insertions(+), 64 deletions(-) diff --git a/packages/clickzetta-sdk/src/sql/cancel.ts b/packages/clickzetta-sdk/src/sql/cancel.ts index cfa3c4436..f0309e172 100644 --- a/packages/clickzetta-sdk/src/sql/cancel.ts +++ b/packages/clickzetta-sdk/src/sql/cancel.ts @@ -13,7 +13,8 @@ export async function cancelJob(opts: ClientOptions, jobId: JobID): Promise 0 && value === undefined && raw.code === undefined) { - throw new ClickZettaApiError("INVALID_CANCEL_RESPONSE", "Missing cancellation status") - } const status = value as Record | undefined const code = status?.errorCode ?? status?.error_code const message = status?.errorMsg ?? status?.error_msg diff --git a/packages/clickzetta-sdk/test/fixtures/cancellation.ts b/packages/clickzetta-sdk/test/fixtures/cancellation.ts index 44f080c21..14d70bbae 100644 --- a/packages/clickzetta-sdk/test/fixtures/cancellation.ts +++ b/packages/clickzetta-sdk/test/fixtures/cancellation.ts @@ -19,6 +19,18 @@ for (const body of [ }) } +// Proto3 omits an empty respStatus even when other fields are populated. +for (const body of [{}, { jobId: "owned-job" }, { requestId: "r1", respStatus: {} }, { code: 0, data: {} }]) { + test(`accepts cancellation without an error status: ${JSON.stringify(body)}`, async () => { + const server = Bun.serve({ port: 0, fetch: () => Response.json(body) }) + try { + await cancelJob({ baseUrl: server.url.origin, tokens: anonymous() }, job) + } finally { + await server.stop(true) + } + }) +} + test("retries cancellation until a late submission becomes terminal", async () => { let attempts = 0 const server = Bun.serve({ diff --git a/packages/cz-cli/src/commands/exec.ts b/packages/cz-cli/src/commands/exec.ts index 9afeacc1b..961219d80 100644 --- a/packages/cz-cli/src/commands/exec.ts +++ b/packages/cz-cli/src/commands/exec.ts @@ -132,14 +132,14 @@ export async function execSql( const timezone = opts?.hints?.["cz.sql.timezone"] const jobId = newJobId(ctx.config.workspace, ctx.instanceId()) opts?.onJobId?.(jobId.id) - const timeoutMs = opts?.timeoutMs ?? ( - opts?.hints?.["sdk.job.timeout"] !== undefined - ? Number(opts.hints["sdk.job.timeout"]) * 1000 - : process.env[CLEANUP_ENV] ? 300_000 : undefined + const explicitTimeoutMs = opts?.timeoutMs ?? ( + opts?.hints?.["sdk.job.timeout"] !== undefined ? Number(opts.hints["sdk.job.timeout"]) * 1000 : undefined ) - if (timeoutMs !== undefined && (!Number.isFinite(timeoutMs) || timeoutMs <= 0)) { + if (explicitTimeoutMs !== undefined && (!Number.isFinite(explicitTimeoutMs) || explicitTimeoutMs <= 0)) { throw new Error("SQL --timeout must be a positive, finite number of seconds") } + // Supervised jobs need a finite fallback in case both processes die before cleanup. + const timeoutMs = explicitTimeoutMs ?? (process.env[CLEANUP_ENV] ? 300_000 : undefined) const lease = await trackSqlJob(ctx.clientOpts, jobId, timeoutMs) const client = { ...ctx.clientOpts, signal: lease.signal } let disposition: "terminal" | "detached" | "cancel" = "cancel" @@ -155,7 +155,8 @@ export async function execSql( jobId, hints: buildExecHints(opts?.hints, traceContext), asynchronous: opts?.asynchronous, - jobTimeoutMs: opts?.asynchronous ? timeoutMs : lease.timeoutMs, + // A detached job outlives its supervisor by design: only an explicit timeout reaches the server. + jobTimeoutMs: opts?.asynchronous ? explicitTimeoutMs : timeoutMs, configStatements: opts?.configStatements, traceparent: traceContext.traceparent, maxRetries: submitMaxRetries(opts?.hints), @@ -193,17 +194,25 @@ export async function execSql( return result } catch (error) { - if (lease.signal.aborted) { - const failure = new Error(`Job ${jobId.id} interrupted or timed out`) - Object.assign(failure, { jobId: jobId.id }) - throw failure - } + if (lease.signal.aborted) throw abortFailure(jobId.id, lease.signal.reason, error) throw error } finally { await lease.finish(disposition) } } +/** Name the abort cause; classifyExecError maps "timed out" to JOB_TIMEOUT and other codes verbatim. */ +function abortFailure(jobId: string, reason: unknown, cause: unknown) { + if (reason instanceof DOMException && reason.name === "TimeoutError") + return Object.assign(new Error(`Job ${jobId} timed out`, { cause }), { jobId }) + const message = reason instanceof Error ? reason.message : "SQL execution interrupted" + const code = (reason as { code?: unknown })?.code + return Object.assign(new Error(`Job ${jobId}: ${message}`, { cause }), { + jobId, + code: typeof code === "string" ? code : "ABORTED", + }) +} + function isAuthError(err: unknown): boolean { if (err instanceof ClickZettaApiError && err.statusCode === 401) return true if (err instanceof Error && err.message.includes("401")) return true diff --git a/packages/cz-cli/src/commands/sql-lifecycle.ts b/packages/cz-cli/src/commands/sql-lifecycle.ts index 824a308a6..c911924c1 100644 --- a/packages/cz-cli/src/commands/sql-lifecycle.ts +++ b/packages/cz-cli/src/commands/sql-lifecycle.ts @@ -3,6 +3,11 @@ import { registerSqlCleanup } from "../sql/cleanup-scope.js" import { parseOutputArgs, renderErrorOutput } from "../output/index.js" const active = new Map }>() +// Abort reasons stay distinct so callers can tell a signal from a lost supervisor +// (the deadline aborts with a TimeoutError DOMException). +const interrupted = () => Object.assign(new Error("SQL execution interrupted"), { code: "ABORTED" }) +const supervisorLost = () => + Object.assign(new Error("SQL cleanup supervisor connection lost"), { code: "SQL_SUPERVISOR_LOST" }) let shutdown: Promise | undefined export function hasActiveSqlJobs() { @@ -18,7 +23,7 @@ export async function trackSqlJob(opts: ClientOptions, job: JobID, timeoutMs?: n let supervisor: Awaited> let completion: Promise | undefined const lease = { - abort: () => controller.abort(new Error("SQL execution interrupted")), + abort: () => controller.abort(interrupted()), finish(state: "terminal" | "detached" | "cancel") { if (completion) return completion completion = (async () => { @@ -65,11 +70,12 @@ export async function trackSqlJob(opts: ClientOptions, job: JobID, timeoutMs?: n } active.set(job.id, lease) try { - supervisor = await registerSqlCleanup(opts, job, signal, timeoutMs ?? 300_000, lease.abort) + supervisor = await registerSqlCleanup(opts, job, signal, timeoutMs ?? 300_000, () => + controller.abort(supervisorLost()), + ) signal.throwIfAborted() return { signal, - timeoutMs, async finish(state: "terminal" | "detached" | "cancel") { await lease.finish(state) // Let the signal owner emit the interruption and choose the exit code. @@ -101,8 +107,13 @@ function stop(signal: string, exitCode: number) { process.stdout.write( renderErrorOutput( { - error: { code: "ABORTED", message: `Execution interrupted by ${signal}.` }, - job_ids: jobs.map(([id]) => id), + error: { + code: "ABORTED", + message: signal === "SIGINT" ? "Execution interrupted by user." : `Execution interrupted by ${signal}.`, + }, + // job_id keeps the pre-existing single-job shape; concurrent execSql callers add job_ids. + ...(jobs[0] && { job_id: jobs[0][0] }), + ...(jobs.length > 1 && { job_ids: jobs.map(([id]) => id) }), }, output.format, output.field, diff --git a/packages/cz-cli/src/commands/sql.ts b/packages/cz-cli/src/commands/sql.ts index dabd3a548..f042459d4 100644 --- a/packages/cz-cli/src/commands/sql.ts +++ b/packages/cz-cli/src/commands/sql.ts @@ -36,6 +36,8 @@ interface SqlArgs extends GlobalArgs { sync: boolean async: boolean timeout: number + /** --timeout as given; undefined when the 300s default applies. */ + explicitTimeout?: number variable?: string[] set?: string[] "job-profile"?: string @@ -379,7 +381,8 @@ async function executeSingle( if (!argv.sync || argv.async) { const asyncHints = { ...hints } - if (argv.timeout) asyncHints["sdk.job.timeout"] = String(argv.timeout) + // The default 300s bounds waiting, not a detached job: only an explicit value is sent. + if (argv.explicitTimeout !== undefined) asyncHints["sdk.job.timeout"] = String(argv.explicitTimeout) const r = await execSqlWithRetry(ctx, sql, { hints: asyncHints, asynchronous: true, configStatements, onJobId }) logOperation("sql", { sql, ok: true, timeMs: Date.now() - t0 }) if (isQueryResult(r)) { @@ -602,10 +605,6 @@ async function emitResult( async function handler(argv: SqlArgs): Promise { const format = argv.format - if (!Number.isFinite(argv.timeout) || argv.timeout <= 0) { - error("USAGE_ERROR", "--timeout must be a positive, finite number of seconds (0 is not supported).", { format, exitCode: 2 }) - return - } if (argv["job-profile"]) { const ctx = await getExecContext(argv) @@ -633,6 +632,10 @@ async function handler(argv: SqlArgs): Promise { } return } + if (!Number.isFinite(argv.timeout) || argv.timeout <= 0) { + error("USAGE_ERROR", "--timeout must be a positive, finite number of seconds (0 is not supported).", { format, exitCode: 2 }) + return + } let sql = resolveSql(argv) if (argv.variable && argv.variable.length > 0) { @@ -853,7 +856,7 @@ export function registerSqlCommand(cli: Argv): void { .option("stdin", { type: "boolean", default: false, describe: "Read SQL from stdin" }) .option("sync", { type: "boolean", default: true, describe: "Wait for query result before returning (default). Use --no-sync or --async for large queries that may take a long time." }) .option("async", { type: "boolean", default: false, describe: "Return job_id immediately without waiting for results. Use for large/long-running queries." }) - .option("timeout", { type: "number", default: 300, describe: "Positive job timeout in seconds (default: 300; 0 is not supported)" }) + .option("timeout", { type: "number", describe: "Positive job timeout in seconds (default: 300; 0 is not supported)" }) .option("variable", { type: "string", array: true, nargs: 1, describe: "Variable substitution: --variable KEY=VALUE. Use ${KEY} in SQL." }) .option("set", { type: "string", array: true, nargs: 1, describe: "Query hint: --set KEY=VALUE (e.g. --set cz.sql.timezone=UTC)" }) .option("job-profile", { type: "string", describe: "Fetch execution profile for a completed job ID (separate from running SQL)" }) @@ -878,7 +881,10 @@ export function registerSqlCommand(cli: Argv): void { `statement itself contains ';', put this in ${CZ_CONFIG_FILE}:`, ` { "${SQL_SPLIT_CONFIG_KEY}": false }`, ].join("\n")), - (argv) => handler(argv as unknown as SqlArgs), + (argv) => { + const args = argv as unknown as SqlArgs + return handler({ ...args, timeout: args.timeout ?? 300, explicitTimeout: args.timeout }) + }, ), ) } diff --git a/packages/cz-cli/src/sql/cleanup-scope.ts b/packages/cz-cli/src/sql/cleanup-scope.ts index 858d19a9d..6f87948e5 100644 --- a/packages/cz-cli/src/sql/cleanup-scope.ts +++ b/packages/cz-cli/src/sql/cleanup-scope.ts @@ -1,5 +1,6 @@ import { z } from "zod" import { createServer, createConnection, type Socket } from "node:net" +import { timingSafeEqual } from "node:crypto" import { abortable, abortAfter, @@ -16,7 +17,8 @@ const registration = z.object({ job: z.object({ id: z.string().min(1).max(256), workspace: z.string().min(1), - instanceId: z.number().int().positive(), + // profile add --verify and setup pass 0 before an instance id is resolved. + instanceId: z.number().int().nonnegative(), }), baseUrl: z .string() @@ -94,7 +96,7 @@ export async function createSqlSupervisor(options: { if (!parsed.success) return cleanup(connection) if (!connection.entry) { const result = registration.safeParse(parsed.data) - if (!result.success || result.data.secret !== secret) return cleanup(connection) + if (!result.success || !sameSecret(result.data.secret, secret)) return cleanup(connection) const key = JSON.stringify([result.data.job.instanceId, result.data.job.workspace, result.data.job.id]) if (jobs.has(key)) return cleanup(connection) jobs.add(key) @@ -168,6 +170,12 @@ export async function createSqlSupervisor(options: { } } +function sameSecret(actual: string, expected: string) { + const a = Buffer.from(actual) + const b = Buffer.from(expected) + return a.length === b.length && timingSafeEqual(a, b) +} + function parseMessage(data: string): unknown { try { return JSON.parse(data) @@ -194,9 +202,11 @@ export async function registerSqlCleanup( const url = new URL(address.url) if (url.hostname !== "127.0.0.1" || url.protocol !== "tcp:" || !url.port) throw new Error("Invalid SQL cleanup supervisor") + // A token refresh is a portal round trip; bound it by the query deadline only. + // The 2s budget covers the loopback connect and handshake. + const credential = await abortable(client.tokens.get(), signal) const deadline = abortAfter(2000, signal) try { - const credential = await abortable(client.tokens.get(), deadline.signal) deadline.signal.throwIfAborted() const socket = createConnection({ host: url.hostname, port: Number(url.port) }) const ready = Promise.withResolvers() @@ -215,7 +225,9 @@ export async function registerSqlCleanup( if (done) return done = true abandon() - const error = new Error(`SQL job ${job.id}: cleanup supervisor connection lost`) + const error = Object.assign(new Error(`SQL job ${job.id}: cleanup supervisor connection lost`), { + code: "SQL_SUPERVISOR_LOST", + }) ready.reject(error) released.reject(error) if (acknowledged) onLost() @@ -301,7 +313,6 @@ function readMessages(socket: Socket, receive: (message: unknown) => void) { socket.on("data", (chunk: Buffer) => { buffer = Buffer.concat([buffer, chunk]) while (!socket.destroyed) { - if (socket.writableLength > 64 * 1024) return socket.destroy() const end = buffer.indexOf(10) if (end < 0) break if (end > 64 * 1024) return socket.destroy() diff --git a/packages/cz-cli/src/sql/supervisor-runtime.ts b/packages/cz-cli/src/sql/supervisor-runtime.ts index 5f06533c0..b08b080a3 100644 --- a/packages/cz-cli/src/sql/supervisor-runtime.ts +++ b/packages/cz-cli/src/sql/supervisor-runtime.ts @@ -1,4 +1,4 @@ -import { appendFile, mkdir } from "node:fs/promises" +import { appendFile, chmod, mkdir } from "node:fs/promises" import { homedir } from "node:os" import path from "node:path" import { CLEANUP_ENV, createSqlSupervisor } from "./cleanup-scope.js" @@ -8,18 +8,21 @@ export async function withSqlSupervisor(run: () => Promise): Promise { const previous = process.env[CLEANUP_ENV] const directory = path.join(process.env.CLICKZETTA_TEST_HOME || homedir(), ".clickzetta") const logfile = path.join(directory, "sql-cleanup.jsonl") - const supervisor = await (async () => { - // Open diagnostics before admitting jobs, not in a failing finalizer. - await mkdir(directory, { recursive: true }) - await appendFile(logfile, "", { mode: 0o600 }) - return createSqlSupervisor({ - async onWarning(warning) { - // Never write into a terminal owned by the TUI renderer or persist credentials. - await appendFile(logfile, JSON.stringify({ time: new Date().toISOString(), ...warning }) + "\n") - }, - }) - })().catch(() => undefined) - // Keep non-SQL tools usable if local sockets/diagnostics are unavailable. + // Diagnostics are best effort: an unwritable ~/.clickzetta loses unconfirmed-cleanup + // records, but must not withdraw supervision and refuse every agent SQL call. + const diagnostics = await mkdir(directory, { recursive: true }) + .then(() => appendFile(logfile, "", { mode: 0o600 })) + // mode only applies on creation; tighten a pre-existing file too. + .then(() => chmod(logfile, 0o600)) + .then(() => true, () => false) + const supervisor = await createSqlSupervisor({ + async onWarning(warning) { + // Never write into a terminal owned by the TUI renderer or persist credentials. + if (!diagnostics) return + await appendFile(logfile, JSON.stringify({ time: new Date().toISOString(), ...warning }) + "\n") + }, + }).catch(() => undefined) + // Keep non-SQL tools usable if the loopback socket is unavailable. // SQL must still fail closed rather than silently run without supervision. process.env[CLEANUP_ENV] = supervisor?.env[CLEANUP_ENV] ?? "unavailable" try { @@ -27,6 +30,7 @@ export async function withSqlSupervisor(run: () => Promise): Promise { } finally { if (previous === undefined) delete process.env[CLEANUP_ENV] else process.env[CLEANUP_ENV] = previous - await supervisor?.close() + // A listener close error has nothing left to report; never replace run()'s exit code. + await supervisor?.close().catch(() => {}) } } diff --git a/packages/cz-cli/test/sql-cancellation.test.ts b/packages/cz-cli/test/sql-cancellation.test.ts index 071db99ed..a7eb86a53 100644 --- a/packages/cz-cli/test/sql-cancellation.test.ts +++ b/packages/cz-cli/test/sql-cancellation.test.ts @@ -35,7 +35,7 @@ test("query deadline aborts a blocked submit and confirms cancellation", async ( expect(await child.exited).not.toBe(0) expect(submitted).toHaveLength(1) expect(cancelled).toEqual(submitted) - expect(await new Response(child.stderr).text()).toContain("interrupted or timed out") + expect(await new Response(child.stderr).text()).toContain(`Job ${submitted[0]} timed out`) } finally { child.kill("SIGKILL") await child.exited @@ -101,7 +101,14 @@ for (const signal of ["SIGTERM", "SIGINT"] as const) { expect(await child.exited).toBe(signal === "SIGTERM" ? 143 : 130) expect(cancelled.sort()).toEqual(submitted.sort()) expect(cancelled).toHaveLength(2) - expect(await new Response(child.stdout).text()).toContain("ABORTED") + const output = JSON.parse(await new Response(child.stdout).text()) + expect(output.error).toEqual({ + code: "ABORTED", + message: signal === "SIGINT" ? "Execution interrupted by user." : "Execution interrupted by SIGTERM.", + }) + // Single-job job_id stays for existing readers; concurrent jobs also list job_ids. + expect(submitted).toContain(output.job_id) + expect(output.job_ids.sort()).toEqual(submitted.sort()) } finally { child.kill("SIGKILL") server.stop(true) diff --git a/packages/cz-cli/test/sql-cleanup-scope.test.ts b/packages/cz-cli/test/sql-cleanup-scope.test.ts index b909c8a4f..e04283be0 100644 --- a/packages/cz-cli/test/sql-cleanup-scope.test.ts +++ b/packages/cz-cli/test/sql-cleanup-scope.test.ts @@ -10,11 +10,11 @@ import { pathToFileURL } from "node:url" import { createSqlSupervisor } from "../src/sql/cleanup-scope.js" const execModule = new URL("../src/commands/exec.ts", import.meta.url).pathname -function program(url: string, options = "{ timeoutMs: 30000 }") { +function program(url: string, options = "{ timeoutMs: 30000 }", instanceId = 1) { return `import { execSql } from ${JSON.stringify(execModule)}; import { anonymous } from '@clickzetta/sdk'; const ctx = { config: { workspace: 'ws', schema: 'public', vcluster: 'vc', instance: 'inst' }, - clientOpts: { baseUrl: ${JSON.stringify(url)}, tokens: anonymous() }, instanceId: () => 1 }; + clientOpts: { baseUrl: ${JSON.stringify(url)}, tokens: anonymous() }, instanceId: () => ${instanceId} }; await execSql(ctx, 'select 1', ${options});` } @@ -408,8 +408,8 @@ test("existing cz Worker environment bridge carries the supervisor endpoint", as } }) -test("supervisor startup failure leaves non-SQL work usable but blocks SQL admission", async () => { - const directory = await mkdtemp(path.join(os.tmpdir(), "cz-supervisor-unavailable-")) +test("unwritable diagnostics keep supervision and SQL admission", async () => { + const directory = await mkdtemp(path.join(os.tmpdir(), "cz-supervisor-diagnostics-")) await Bun.write(path.join(directory, ".clickzetta"), "not a directory") const remote = queryServer() const child = Bun.spawn( @@ -419,11 +419,11 @@ test("supervisor startup failure leaves non-SQL work usable but blocks SQL admis ` import { withSqlSupervisor } from ${JSON.stringify(new URL("../src/sql/supervisor-runtime.ts", import.meta.url).pathname)}; await withSqlSupervisor(async () => { - console.log('non-sql-work-ran'); + if (!process.env.CZ_SQL_CLEANUP?.startsWith('{')) throw new Error('supervision disabled'); const { execSql } = await import(${JSON.stringify(execModule)}); const { anonymous } = await import('@clickzetta/sdk'); await execSql({ config: { workspace: 'ws', schema: 'public', vcluster: 'vc', instance: 'inst' }, - clientOpts: { baseUrl: ${JSON.stringify(remote.server.url.origin)}, tokens: anonymous() }, instanceId: () => 1 }, 'select 1'); + clientOpts: { baseUrl: ${JSON.stringify(remote.server.url.origin)}, tokens: anonymous() }, instanceId: () => 1 }, 'select 1', { timeoutMs: 300 }); }); `, ], @@ -431,9 +431,12 @@ test("supervisor startup failure leaves non-SQL work usable but blocks SQL admis ) try { expect(await child.exited).not.toBe(0) - expect(await new Response(child.stdout).text()).toContain("non-sql-work-ran") - expect(await new Response(child.stderr).text()).toContain("SQL cleanup supervisor unavailable") - expect(remote.submitted).toEqual([]) + const stderr = await new Response(child.stderr).text() + expect(stderr).toContain("timed out") + expect(stderr).not.toContain("unavailable") + expect(remote.submitted).toHaveLength(1) + // Child and supervisor may both cancel; either confirms the timed-out job. + expect(remote.cancelled).toContain(remote.submitted[0]) } finally { child.kill("SIGKILL") await child.exited @@ -441,3 +444,24 @@ test("supervisor startup failure leaves non-SQL work usable but blocks SQL admis await rm(directory, { recursive: true, force: true }) } }) + +// profile add --verify and setup submit before any instance id is resolved. +test("instance id 0 registers under supervision", async () => { + const remote = queryServer() + const supervisor = await createSqlSupervisor({ onWarning: async () => {} }) + const child = Bun.spawn([process.execPath, "--eval", program(remote.server.url.origin, "{ timeoutMs: 300 }", 0)], { + env: { ...process.env, ...supervisor.env }, + stdout: "pipe", + stderr: "pipe", + }) + try { + await child.exited + expect(await new Response(child.stderr).text()).not.toContain("supervisor connection lost") + expect(remote.submitted).toHaveLength(1) + } finally { + child.kill("SIGKILL") + await child.exited + await supervisor.close() + await remote.server.stop(true) + } +}) diff --git a/packages/cz-cli/test/sql-timeout-options.test.ts b/packages/cz-cli/test/sql-timeout-options.test.ts index f0978caa3..e25393386 100644 --- a/packages/cz-cli/test/sql-timeout-options.test.ts +++ b/packages/cz-cli/test/sql-timeout-options.test.ts @@ -43,3 +43,15 @@ for (const command of ["table list", "schema list"]) { expect(timeouts.every((value) => value === undefined)).toBe(true) }) } + +test("sql --async without --timeout leaves the detached job on the deployment timeout", async () => { + const result = await execute("sql", ["select 1", "--async"]) + expect(result.exitCode).toBe(0) + expect(timeouts).toEqual([undefined]) +}) + +test("sql --job-profile ignores an unused --timeout", async () => { + onFetch({ match: (url) => url.includes("/lh/getJob"), respond: () => Response.json({ status: { state: "SUCCEED" } }) }) + const result = await execute("sql", ["--job-profile", "job1", "--timeout", "0"]) + expect(result.exitCode).toBe(0) +}) diff --git a/specs/sql-cancellation.md b/specs/sql-cancellation.md index 9d7be345e..86b4dd074 100644 --- a/specs/sql-cancellation.md +++ b/specs/sql-cancellation.md @@ -19,9 +19,10 @@ process runner. A separate connection owns each SQL job, so concurrent commands and sessions do not share job ownership. Remote servers use the supervisor on the machine actually executing their CLI subprocesses, not the attached client. -If the supervisor or its diagnostic destination cannot start, non-SQL tools can -still run. The environment explicitly marks supervision unavailable and SQL -admission fails closed. A missing inherited environment (for example after an +If the supervisor cannot start, non-SQL tools can still run. The environment +explicitly marks supervision unavailable and SQL admission fails closed. An +unwritable diagnostic destination only drops unconfirmed-cleanup records; it +does not withdraw supervision. A missing inherited environment (for example after an explicit `env -i`) uses standalone CLI behavior; it cannot promise supervisor recovery. @@ -54,7 +55,8 @@ handoff. Disconnect after a successful handoff does not cancel the detached job. | Boundary | Budget | | ----------------------- | --------------------------------------------- | -| Registration | 2 seconds, also subject to the query deadline | +| Credential resolution | Query deadline only (may refresh a token) | +| Registration handshake | 2 seconds, also subject to the query deadline | | Child cleanup | 1.5 seconds per job, concurrently | | Signal shutdown | Hard 2-second exit bound | | Supervisor cleanup | 5 seconds per job, concurrently | @@ -72,13 +74,21 @@ supervisor and attempts cancellation of all outstanding registrations. Standalone `cz-cli sql` retains its default 300-second timeout and now actively cancels on timeout, SIGINT or SIGTERM. Positive explicit timeouts are sent to the -server, including for `--async`. Invalid/non-positive CLI timeout values fail at -the command boundary; zero is not an unlimited-timeout sentinel. +server, including for `--async`; the 300-second default is not, because it bounds +waiting rather than a detached job. Invalid/non-positive CLI timeout values fail +at the command boundary (except `--job-profile`, which submits nothing); zero is +not an unlimited-timeout sentinel. + +Interruption errors name their cause: a deadline is `Job timed out` +(`JOB_TIMEOUT`), a signal is `ABORTED`, a lost supervisor is +`SQL_SUPERVISOR_LOST`. The signal envelope keeps `job_id` and adds `job_ids` +when several jobs were active. `cancelJob` treats a 2xx without a populated error +status as accepted; confirmation still comes from polling job state. Standalone `execSql` callers without a timeout (such as table/schema/file commands) keep their previous deployment-defined timeout. They still clean up on caught failure and signals. Agent-supervised calls without an explicit timeout receive -a 300-second fallback. No upstream hook exposes the shell's remaining budget, so +a 300-second fallback (not sent to the server for `--async` handoff). No upstream hook exposes the shell's remaining budget, so supervision uses the SQL deadline plus disconnect/heartbeat detection instead. If the entire agent is killed, loses network access, or the credential snapshot From 13f661bf9e8aba5738cce1f1cf86424e3c5ea475 Mon Sep 17 00:00:00 2001 From: suibianwanwan Date: Thu, 8 Oct 2026 11:50:45 +0800 Subject: [PATCH 5/5] fix(sql): skip cancellation for rejected submissions A definitive 4xx from submitJob means no job was created, so waiting 1.5s to cancel it only delayed batch failures and printed a false unconfirmed-cancellation warning. Co-Authored-By: Claude Opus 5.5 (1M context) --- packages/cz-cli/src/commands/exec.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/cz-cli/src/commands/exec.ts b/packages/cz-cli/src/commands/exec.ts index 961219d80..bde7ef228 100644 --- a/packages/cz-cli/src/commands/exec.ts +++ b/packages/cz-cli/src/commands/exec.ts @@ -160,6 +160,12 @@ export async function execSql( configStatements: opts?.configStatements, traceparent: traceContext.traceparent, maxRetries: submitMaxRetries(opts?.hints), + }).catch((error: unknown) => { + // The gateway answered with a definitive client rejection: no job exists to cancel. + // Timeouts, conflicts and throttling stay ambiguous and still cancel. + const status = (error as { statusCode?: number })?.statusCode + if (status && status >= 400 && status < 500 && ![408, 409, 429].includes(status)) disposition = "terminal" + throw error }) if (opts?.asynchronous) { disposition = "detached"