fix: handle quota and billing errors

Co-Authored-By: cz-cli <noreply@clickzetta.com>

Author: suibianwanwan

packages/cz-cli/src/commands/billing-error.ts

@@ -0,0 +1,55 @@
+import { accountLoginUrlForService } from "./account-login.js"
+import { getDefaultProfileName, loadProfiles } from "../connection/profile-store.js"
+
+const INSUFFICIENT_BALANCE_CODE = "CZLH-60029"
+const OVERDUE_PAYMENTS_RE = /Account\s+([a-z0-9_-]+)\s+has overdue payments\./i
+const INSUFFICIENT_BALANCE_RE = /insufficient account balance|overdue payments|job submission is currently restricted/i
+
+function activeProfileEntry(profileName?: string) {
+  const profiles = loadProfiles()
+  const explicitName = profileName ?? process.env.CZ_PROFILE
+  if (explicitName) return profiles[explicitName]
+  return profiles[getDefaultProfileName() ?? ""] ?? Object.values(profiles)[0]
+}
+
+function normalizeAccountsUrl(value: string) {
+  return value.trim().replace(/\/+$/, "")
+}
+
+function extractAccountName(message: string) {
+  return OVERDUE_PAYMENTS_RE.exec(message)?.[1]
+}
+
+function resolveAccountsUrl(input: {
+  accountName?: string
+  profileName?: string
+  service?: string
+}) {
+  const profile = activeProfileEntry(input.profileName)
+  if (typeof profile?.accounts_url === "string" && profile.accounts_url.trim()) {
+    return normalizeAccountsUrl(profile.accounts_url)
+  }
+  if (!input.accountName || !input.service) return undefined
+  return accountLoginUrlForService(input.service, input.accountName)
+}
+
+export function formatBillingError(input: {
+  code?: string
+  message?: string
+  profileName?: string
+  service?: string
+}) {
+  const message = input.message ?? "Query failed"
+  if (input.code !== INSUFFICIENT_BALANCE_CODE && !INSUFFICIENT_BALANCE_RE.test(message)) {
+    return message
+  }
+
+  const accountsUrl = resolveAccountsUrl({
+    accountName: extractAccountName(message),
+    profileName: input.profileName,
+    service: input.service,
+  })
+  if (!accountsUrl) return message
+
+  return `Insufficient account balance. Please visit ${accountsUrl} to add funds.`
+}

packages/cz-cli/src/commands/sql.ts

@@ -6,6 +6,7 @@ import { success, successRows, error, parseOutputArgs, renderOutput } from "../o
 import { maskRows } from "../output/masking.js"
 import { logOperation } from "../logger.js"
 import { getExecContext, execSql, execSqlWithRetry, isQueryResult, validateIdentifier, classifyExecError, type ExecContext } from "./exec.js"
+import { formatBillingError } from "./billing-error.js"
 
 const WRITE_RE = /^\s*(INSERT|UPDATE|DELETE|REPLACE|ALTER|CREATE|DROP|TRUNCATE|RENAME|FORK)\b/i
 const SELECT_RE = /^\s*(SELECT\b|WITH\b[\s\S]*?\bSELECT\b|SHOW\b)/i
@@ -131,6 +132,7 @@ async function applyUseStatement(
   hints?: Record<string, string>,
   configStatements?: string[],
   timeoutMs?: number,
+  profileName?: string,
 ): Promise<boolean> {
   if (use.kind === "workspace") {
     ctx.config.workspace = use.target
@@ -145,7 +147,7 @@ async function applyUseStatement(
     }
     if (result.status === JobStatus.FAILED) {
       logOperation("sql", { sql: use.normalized, ok: false, errorCode: result.errorCode })
-      error(result.errorCode ?? "SCHEMA_NOT_FOUND", result.errorMessage ?? `Schema '${use.target}' does not exist.`, { format })
+      error(result.errorCode ?? "SCHEMA_NOT_FOUND", formatQueryError(result, ctx, profileName, `Schema '${use.target}' does not exist.`), { format })
       return false
     }
     ctx.config.schema = extractSchema(use.raw)
@@ -159,7 +161,7 @@ async function applyUseStatement(
   }
   if (result.status === JobStatus.FAILED) {
     logOperation("sql", { sql: use.normalized, ok: false, errorCode: result.errorCode })
-    error(result.errorCode ?? "SQL_ERROR", result.errorMessage ?? "Query failed", { format })
+    error(result.errorCode ?? "SQL_ERROR", formatQueryError(result, ctx, profileName), { format })
     return false
   }
   if (!useTargetExists(result.rows, use.target)) {
@@ -241,7 +243,7 @@ async function executeSingle(
   // Intercept USE statements — client-side context switch
   const use = parseUseStatement(sql)
   if (use) {
-    if (!await applyUseStatement(ctx, use, format, hints, configStatements, argv.timeout * 1000)) return
+    if (!await applyUseStatement(ctx, use, format, hints, configStatements, argv.timeout * 1000, argv.profile)) return
     success({ use: use.normalized }, { format, timeMs: 0 })
     return
   }
@@ -296,7 +298,7 @@ async function executeSingle(
       r = await execSqlWithRetry(ctx, sql, { hints: probeHints, timeoutMs: argv.timeout * 1000, configStatements, onJobId })
       if (!isQueryResult(r)) { error("UNEXPECTED_RESULT", "Expected query result but got async marker.", { format }); return }
       if (r.status === JobStatus.FAILED) {
-        await handleFailure(r, sql, ctx, format, t0)
+        await handleFailure(r, sql, ctx, format, t0, argv.profile)
         return
       }
       await emitResult(r, sql, argv, ctx, t0)
@@ -306,7 +308,7 @@ async function executeSingle(
     if (r.status === JobStatus.FAILED) {
       const hint = await fetchSchemaHint(ctx, sql, r.errorMessage ?? "")
       logOperation("sql", { sql, ok: false, errorCode: r.errorCode, timeMs: Date.now() - t0 })
-      error(r.errorCode ?? "SQL_ERROR", r.errorMessage ?? "Query failed", { format, extra: hint ? { schema: hint } : undefined })
+      error(r.errorCode ?? "SQL_ERROR", formatQueryError(r, ctx, argv.profile), { format, extra: hint ? { schema: hint } : undefined })
       return
     }
     if (r.rowCount > rowLimit) {
@@ -342,7 +344,7 @@ async function executeSingle(
       const r = await execSqlWithRetry(ctx, probeSql, { hints, timeoutMs: argv.timeout * 1000, configStatements, onJobId })
       if (!isQueryResult(r)) { error("UNEXPECTED_RESULT", "Expected query result but got async marker.", { format }); return }
       if (r.status === JobStatus.FAILED) {
-        await handleFailure(r, sql, ctx, format, t0)
+        await handleFailure(r, sql, ctx, format, t0, argv.profile)
         return
       }
       let aiMessage: string | undefined
@@ -360,20 +362,29 @@ async function executeSingle(
   const r = await execSqlWithRetry(ctx, sql, { hints, timeoutMs: argv.timeout * 1000, configStatements, onJobId })
   if (!isQueryResult(r)) { error("UNEXPECTED_RESULT", "Expected query result but got async marker.", { format }); return }
   if (r.status === JobStatus.FAILED) {
-    await handleFailure(r, sql, ctx, format, t0)
+    await handleFailure(r, sql, ctx, format, t0, argv.profile)
     return
   }
 
   await emitResult(r, sql, argv, ctx, t0)
 }
 
-async function handleFailure(r: QueryResult, sql: string, ctx: ExecContext, format: string, t0: number): Promise<void> {
+function formatQueryError(r: QueryResult, ctx: ExecContext, profileName?: string, fallback = "Query failed") {
+  return formatBillingError({
+    code: r.errorCode,
+    message: r.errorMessage ?? fallback,
+    profileName,
+    service: ctx.config.service,
+  })
+}
+
+async function handleFailure(r: QueryResult, sql: string, ctx: ExecContext, format: string, t0: number, profileName?: string): Promise<void> {
   const hint = await fetchSchemaHint(ctx, sql, r.errorMessage ?? "")
   logOperation("sql", { sql, ok: false, errorCode: r.errorCode, timeMs: Date.now() - t0 })
   const aiMessage = hint
     ? `SQL failed. Available schema info attached in the 'schema' field — check table/column names and retry.`
     : undefined
-  error(r.errorCode ?? "SQL_ERROR", r.errorMessage ?? "Query failed", {
+  error(r.errorCode ?? "SQL_ERROR", formatQueryError(r, ctx, profileName), {
     format,
     extra: hint ? { schema: hint } : undefined,
     ...(aiMessage && { aiMessage }),
@@ -498,7 +509,7 @@ async function handler(argv: SqlArgs): Promise<void> {
         // Extract USE statements to update session context client-side
         const use = parseUseStatement(stmt)
         if (use) {
-          if (!await applyUseStatement(ctx, use, format, accumulatedHints, configStatements, argv.timeout * 1000)) return
+          if (!await applyUseStatement(ctx, use, format, accumulatedHints, configStatements, argv.timeout * 1000, argv.profile)) return
           configStatements.push(stmt)
           continue
         }
@@ -508,7 +519,7 @@ async function handler(argv: SqlArgs): Promise<void> {
           try {
             const r = await execSqlWithRetry(ctx, stmt, { hints: accumulatedHints, timeoutMs: argv.timeout * 1000, configStatements })
             if (isQueryResult(r) && r.status === JobStatus.FAILED) {
-              const line = { index: i, sql: stmt, error: { code: r.errorCode ?? "SQL_ERROR", message: r.errorMessage ?? "Query failed" }, time_ms: Date.now() - t0, ...(r.jobId ? { job_id: r.jobId } : {}) }
+              const line = { index: i, sql: stmt, error: { code: r.errorCode ?? "SQL_ERROR", message: formatQueryError(r, ctx, argv.profile) }, time_ms: Date.now() - t0, ...(r.jobId ? { job_id: r.jobId } : {}) }
               process.stdout.write((format === "pretty" ? renderOutput(line, format) : JSON.stringify(line)) + "\n")
               logOperation("sql", { sql: stmt, ok: false, errorCode: r.errorCode })
             } else if (isQueryResult(r)) {
@@ -529,7 +540,7 @@ async function handler(argv: SqlArgs): Promise<void> {
           const r = await execSqlWithRetry(ctx, stmt, { hints: accumulatedHints, timeoutMs: argv.timeout * 1000, configStatements })
           if (isQueryResult(r) && r.status === JobStatus.FAILED) {
             logOperation("sql", { sql: stmt, ok: false, errorCode: r.errorCode })
-            error(r.errorCode ?? "SQL_ERROR", r.errorMessage ?? "Query failed", { format })
+            error(r.errorCode ?? "SQL_ERROR", formatQueryError(r, ctx, argv.profile), { format })
             return
           }
         } else {

packages/cz-cli/src/llm/clickzetta-rotation.ts

@@ -12,11 +12,13 @@ const API = {
   SAVE: "/llm-gateway-admin/v2/virtual-key/save",
   GET: "/llm-gateway-admin/v2/virtual-key/getApiKey",
 }
-const QUOTA_EXHAUSTED_PATTERN = /virtual key total quota exceeded/i
-const ALIAS_PREFIX = "cz-code_auto_"
-const PROFILE_ALIAS_PREFIX = "cz-cli_auto_"
+const FREE_ALIAS_PREFIX = "cz-code_auto_"
+const AUTO_ALIAS_PREFIX = "cz-cli_auto_"
+const PROFILE_ALIAS_PREFIX = "cz-cli_user_"
 const ALIAS_ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789"
-export const CLICKZETTA_ROTATION_PROMPT = "Free quota exhausted. Create a new virtual key with the current profile and switch?"
+const VIRTUAL_KEY_QUOTA_EXHAUSTED_PATTERN = /virtual key\s+\w+\s+quota exceeded/i
+const VIRTUAL_KEY_NAME_PATTERN = /virtual key\s+['"]([^'"]+)['"]/i
+export const CLICKZETTA_ROTATION_PROMPT = "Your complimentary token quota has been exhausted.\nWe also offer competitively priced paid token plans, and I'd be happy to help you create and configure a paid API key."
 export const CLICKZETTA_ROTATION_HEADER = "Quota"
 export const CLICKZETTA_ROTATION_CONFIRM_LABEL = "Create & switch"
 export const CLICKZETTA_ROTATION_CANCEL_LABEL = "Keep current"
@@ -85,11 +87,20 @@ export function inferAiGatewayUrl(profile: { service?: string; instance?: string
 
 function randomAlias() {
   const bytes = crypto.getRandomValues(new Uint8Array(8))
-  return ALIAS_PREFIX + Array.from(bytes, (byte) => ALIAS_ALPHABET[byte % ALIAS_ALPHABET.length]).join("")
+  return AUTO_ALIAS_PREFIX + Array.from(bytes, (byte) => ALIAS_ALPHABET[byte % ALIAS_ALPHABET.length]).join("")
 }
 
-function quotaExhausted(detail?: string | null) {
-  return typeof detail === "string" && QUOTA_EXHAUSTED_PATTERN.test(detail)
+export function clickzettaQuotaVirtualKeyName(detail?: string | null) {
+  if (typeof detail !== "string") return undefined
+  return VIRTUAL_KEY_NAME_PATTERN.exec(detail)?.[1]
+}
+
+export function isClickzettaVirtualKeyQuotaErrorDetail(detail?: string | null) {
+  return typeof detail === "string" && VIRTUAL_KEY_QUOTA_EXHAUSTED_PATTERN.test(detail) && !!clickzettaQuotaVirtualKeyName(detail)
+}
+
+export function isClickzettaFreeQuotaErrorDetail(detail?: string | null) {
+  return isClickzettaVirtualKeyQuotaErrorDetail(detail) && clickzettaQuotaVirtualKeyName(detail)?.startsWith(FREE_ALIAS_PREFIX) === true
 }
 
 function promptAllowed(approval: RotateOptions["approval"]) {
@@ -215,7 +226,7 @@ export function isClickzettaQuotaExhausted(input: {
   status?: number
   detail?: string | null
 }) {
-  return input.provider === "clickzetta" && input.status === 429 && quotaExhausted(input.detail)
+  return input.provider === "clickzetta" && input.status === 429 && isClickzettaFreeQuotaErrorDetail(input.detail)
 }
 
 export async function rotateClickzettaLlm(input: {

packages/cz-cli/test/clickzetta-rotation.test.ts

@@ -113,20 +113,21 @@ afterAll(() => {
 })
 
 describe("clickzetta key rotation", () => {
-  test("recognizes the clickzetta quota-exhausted 429 pattern", () => {
+  test("recognizes the clickzetta free key quota-exhausted 429 pattern", () => {
     expect(
       isClickzettaQuotaExhausted({
         provider: "clickzetta",
         status: 429,
         detail:
-          "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-code_auto_old', current usage: 10075371 tokens\"}",
+          "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-code_auto_pdiaxzjq', current usage: 10082801 tokens\"}",
       }),
     ).toBe(true)
     expect(
       isClickzettaQuotaExhausted({
         provider: "clickzetta",
-        status: 401,
-        detail: "Invalid virtual key",
+        status: 429,
+        detail:
+          "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-cli_auto_UAT_TEST', current usage: 10075371 tokens\"}",
       }),
     ).toBe(false)
   })
@@ -147,11 +148,11 @@ describe("clickzetta key rotation", () => {
     expect(studioCalls[0]?.body).toEqual({
       pageIndex: 1,
       pageSize: 200,
-      vApiKeyAlias: "cz-cli_auto_UAT_TEST",
+      vApiKeyAlias: "cz-cli_user_UAT_TEST",
     })
     expect(studioCalls[1]?.path).toBe("/llm-gateway-admin/v2/virtual-key/save")
     expect(studioCalls[1]?.body).toEqual({
-      vApiKeyAlias: "cz-cli_auto_UAT_TEST",
+      vApiKeyAlias: "cz-cli_user_UAT_TEST",
       rateLimitConfigs: { quota_total: 10000000 },
     })
     const profiles = readFileSync(profileFile, "utf-8")
@@ -216,31 +217,60 @@ describe("clickzetta key rotation", () => {
     expect(studioCalls[0]?.path).toBe("/llm-gateway-admin/v2/virtual-key/listWithAuth")
     expect(studioCalls[1]?.path).toBe("/llm-gateway-admin/v2/virtual-key/save")
     expect(studioCalls[1]?.body).toEqual({
-      vApiKeyAlias: expect.stringMatching(/^cz-code_auto_/),
+      vApiKeyAlias: expect.stringMatching(/^cz-cli_auto_/),
       rateLimitConfigs: { quota_total: 10000000 },
     })
     const profiles = readFileSync(profileFile, "utf-8")
     expect(profiles).toContain('api_key = "ck-new"')
   })
 
   test("reuses an existing key with the derived alias before creating a new one", async () => {
-    gatewayState.listData = [{ id: 77, vApiKeyAlias: "cz-cli_auto_UAT_TEST" }]
+    writeFileSync(
+      profileFile,
+      [
+        'default_profile = "uat"',
+        "",
+        "[profiles.uat]",
+        'pat = "pat-uat"',
+        'instance = "inst-uat"',
+        'workspace = "ws-uat"',
+        'username = "EXISTING_TEST"',
+        'service = "uat-service.example"',
+        'protocol = "https"',
+        'ai_gateway_url = "https://mock-aimesh.example/gateway/v1"',
+        "",
+      ].join("\n"),
+    )
+    gatewayState.listData = [{ id: 77, vApiKeyAlias: "cz-cli_user_EXISTING_TEST" }]
     gatewayState.getById[77] = "ck-existing"
 
     const result = await maybeRotateExhaustedClickzettaLlm({
       provider: "clickzetta",
       status: 429,
       detail:
-        "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-cli_auto_UAT_TEST', current usage: 10075371 tokens\"}",
+        "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-code_auto_old', current usage: 10075371 tokens\"}",
       approval: "auto",
     })
 
     expect(result?.rotated).toBe(true)
-    expect(result?.alias).toBe("cz-cli_auto_UAT_TEST")
+    expect(result?.alias).toBe("cz-cli_user_EXISTING_TEST")
     expect(studioCalls.map((call) => call.path)).toEqual([
       "/llm-gateway-admin/v2/virtual-key/listWithAuth",
       "/llm-gateway-admin/v2/virtual-key/getApiKey?id=77",
     ])
     expect(readFileSync(profileFile, "utf-8")).toContain('api_key = "ck-existing"')
   })
+
+  test("does not rotate when quota is exhausted for a non-free key", async () => {
+    const result = await maybeRotateExhaustedClickzettaLlm({
+      provider: "clickzetta",
+      status: 429,
+      detail:
+        "{\"code\":429,\"message\":\"Virtual key total quota exceeded: limit is 10000000 tokens for virtual key 'cz-cli_auto_UAT_TEST', current usage: 10075371 tokens\"}",
+      approval: "auto",
+    })
+
+    expect(result).toBeUndefined()
+    expect(studioCalls).toEqual([])
+  })
 })