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..f0309e172 100644 --- a/packages/clickzetta-sdk/src/sql/cancel.ts +++ b/packages/clickzetta-sdk/src/sql/cancel.ts @@ -1,21 +1,79 @@ -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 individually, so an absent status (alone or beside + // other fields) is success; only a populated error status rejects. + 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") } + 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 +} - const resp = await request(opts, "/lh/cancelJob", body) - return resp +export type CancellationResult = { confirmed: true; state: string } | { confirmed: false; reason: string } + +/** 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..14d70bbae --- /dev/null +++ b/packages/clickzetta-sdk/test/fixtures/cancellation.ts @@ -0,0 +1,138 @@ +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) + } + }) +} + +// 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({ + 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 { + // 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" + return { jobId: jobId.id, status: "RUNNING" as const } + } + // HYBRID mode: submitJob may return the result directly if the query + // finished within hybridPollingTimeout. Check for a terminal state. + const raw = submitResp as { status?: { state?: string } } + let result: QueryResult + if (raw?.status?.state && ["SUCCEED", "FAILED", "CANCELLED"].includes(raw.status.state)) { + const errorCode = retryableSubmitCode(submitResp) + if (isRetryableErrorCode(errorCode)) { + result = await pollJobResult(client, jobId, { jobTimeoutMs: opts?.timeoutMs, timezone }) + } else { + result = await parseJobResponse(submitResp as Parameters[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" + + // 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) throw abortFailure(jobId.id, lease.signal.reason, error) + throw error + } finally { + await lease.finish(disposition) } +} - return result +/** 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 { 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..c911924c1 --- /dev/null +++ b/packages/cz-cli/src/commands/sql-lifecycle.ts @@ -0,0 +1,125 @@ +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 }>() +// 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() { + 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(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, () => + controller.abort(supervisorLost()), + ) + signal.throwIfAborted() + return { + signal, + 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: 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, + ) + "\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..59b9f3535 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)) { @@ -629,25 +632,18 @@ 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) { 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 +741,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 +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: "Job timeout in seconds (default: 300)" }) + .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)" }) @@ -887,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/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..6f87948e5 --- /dev/null +++ b/packages/cz-cli/src/sql/cleanup-scope.ts @@ -0,0 +1,325 @@ +import { z } from "zod" +import { createServer, createConnection, type Socket } from "node:net" +import { timingSafeEqual } from "node:crypto" +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), + // profile add --verify and setup pass 0 before an instance id is resolved. + instanceId: z.number().int().nonnegative(), + }), + 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 || !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) + 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 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) + } 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") + // 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 { + 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 = 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() + } + 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) { + 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..b08b080a3 --- /dev/null +++ b/packages/cz-cli/src/sql/supervisor-runtime.ts @@ -0,0 +1,36 @@ +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" + +/** 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") + // 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 { + return await run() + } finally { + if (previous === undefined) delete process.env[CLEANUP_ENV] + else process.env[CLEANUP_ENV] = previous + // 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 new file mode 100644 index 000000000..a7eb86a53 --- /dev/null +++ b/packages/cz-cli/test/sql-cancellation.test.ts @@ -0,0 +1,178 @@ +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(`Job ${submitted[0]} 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) + 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) + } + }, 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..e04283be0 --- /dev/null +++ b/packages/cz-cli/test/sql-cleanup-scope.test.ts @@ -0,0 +1,467 @@ +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 }", 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: () => ${instanceId} }; + 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("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( + [ + process.execPath, + "--eval", + ` + import { withSqlSupervisor } from ${JSON.stringify(new URL("../src/sql/supervisor-runtime.ts", import.meta.url).pathname)}; + await withSqlSupervisor(async () => { + 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', { timeoutMs: 300 }); + }); + `, + ], + { env: { ...process.env, CLICKZETTA_TEST_HOME: directory }, stdout: "pipe", stderr: "pipe" }, + ) + try { + expect(await child.exited).not.toBe(0) + 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 + await remote.server.stop(true) + 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 new file mode 100644 index 000000000..e25393386 --- /dev/null +++ b/packages/cz-cli/test/sql-timeout-options.test.ts @@ -0,0 +1,57 @@ +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) + }) +} + +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 new file mode 100644 index 000000000..86b4dd074 --- /dev/null +++ b/specs/sql-cancellation.md @@ -0,0 +1,107 @@ +# 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 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. + +## 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 | +| ----------------------- | --------------------------------------------- | +| 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 | +| 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`; 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 (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 +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.