feat: add local auth billing and usage management

lib/server/generation-service.ts

@@ -20,13 +20,14 @@ import {
 import {
   createGenerationJob,
   getGenerationJob,
-  recordUsageEvent,
+  recordUsageForJob,
   updateGenerationJob
 } from "@/lib/server/data-store";
+import { chargeGenerationJob, quoteGenerationCharge } from "@/lib/server/billing-service";
 import { createMockImageBuffer } from "@/lib/server/mock-image";
 import { importRemoteImageAsAsset, saveGeneratedAsset } from "@/lib/server/storage";
 import { DEFAULT_OWNER_ID, toAbsoluteUrl } from "@/lib/server/runtime";
-import type { EnabledImageCapability, GenerationJob, VisualTaskQueryResponse } from "@/lib/types";
+import type { BillingJobCharge, EnabledImageCapability, GenerationJob, UsageContext, VisualTaskQueryResponse } from "@/lib/types";
 import { queryVisualTask, shouldMockVisualApi, submitVisualTask } from "@/lib/volcengine/visual-client";
 import {
   bailianResultUrls,
@@ -53,18 +54,91 @@ export type SubmitImageJobInput = {
   min_ratio?: number;
   max_ratio?: number;
   force_single?: boolean;
-  resolution?: "4k" | "8k";
   quality?: string;
-  seed?: number;
   retryOf?: string;
   idempotencyKey?: string;
   idempotencyFingerprint?: string;
   priority?: number;
   maxAttempts?: number;
   webhookUrl?: string;
+  usageContext?: UsageContext;
 };
 
+type PreparedImageGeneration = {
+  ownerId: string;
+  normalizedUrls: string[];
+  engine: ImageCreationEngine;
+  providerPayload: Record<string, unknown>;
+  missingBailianKey: boolean;
+  provider: "mock" | "bailian" | "evolink" | "volcengine-visual";
+  reqKey: string;
+  requestPayload: Record<string, unknown>;
+};
+
+export async function quoteImageGeneration(input: SubmitImageJobInput, origin: string): Promise<BillingJobCharge | undefined> {
+  const prepared = prepareImageGeneration(input, origin);
+  if (prepared.missingBailianKey) return undefined;
+  return quoteGenerationCharge({
+    provider: prepared.provider,
+    capability: input.capability,
+    reqKey: prepared.reqKey,
+    requestPayload: prepared.requestPayload,
+    usageContext: input.usageContext,
+    externalClientId: input.externalClientId,
+    allowUnboundOrganization: true
+  });
+}
+
 export async function submitImageJob(input: SubmitImageJobInput, origin: string): Promise<GenerationJob> {
+  const prepared = prepareImageGeneration(input, origin);
+  const { ownerId, normalizedUrls, providerPayload, missingBailianKey, provider, reqKey, requestPayload } = prepared;
+  const billing = missingBailianKey ? undefined : await quoteGenerationCharge({
+    provider,
+    capability: input.capability,
+    reqKey,
+    requestPayload,
+    usageContext: input.usageContext,
+    externalClientId: input.externalClientId
+  });
+  let job = await createGenerationJob({
+    ownerId,
+    externalClientId: input.externalClientId,
+    capability: input.capability,
+    provider,
+    reqKey,
+    status: missingBailianKey ? "failed" : "queued",
+    prompt: input.prompt,
+    inputAssetIds: input.inputAssetIds || [],
+    inputUrls: normalizedUrls,
+    outputAssetIds: [],
+    requestPayload,
+    error: missingBailianKey ? { message: "缺少 BAILIAN_API_KEY，请先在设置页配置阿里云百炼 API Key。", retryable: true } : undefined,
+    retryOf: input.retryOf,
+    idempotencyKey: input.idempotencyKey,
+    idempotencyFingerprint: input.idempotencyFingerprint,
+    priority: input.priority,
+    maxAttempts: input.maxAttempts,
+    webhookUrl: input.webhookUrl,
+    usageContext: input.usageContext,
+    billing
+  });
+  if (!billing) return job;
+  try {
+    return await chargeGenerationJob(job);
+  } catch (error) {
+    await updateGenerationJob(job.id, {
+      status: "failed",
+      billing: { ...billing, status: "not_charged" },
+      error: {
+        message: error instanceof Error ? error.message : String(error),
+        retryable: false
+      }
+    }).catch(() => undefined);
+    throw error;
+  }
+}
+
+function prepareImageGeneration(input: SubmitImageJobInput, origin: string): PreparedImageGeneration {
   const ownerId = input.ownerId || DEFAULT_OWNER_ID;
   const capability = getEnabledImageCapability(input.capability);
   const normalizedUrls = (input.imageUrls || []).map((url) => toAbsoluteUrl(url, origin));
@@ -78,38 +152,20 @@ export async function submitImageJob(input: SubmitImageJobInput, origin: string)
   const mock = engine === "bailian" ? shouldMockBailian() : engine === "evolink" ? shouldMockEvolinkApi() : shouldMockVisualApi();
   const missingBailianKey = engine === "bailian" && !mock && !bailianConfig.apiKey;
   const reqKey = engine === "bailian" ? bailianConfig.imageModel : engine === "evolink" ? getEvolinkImageSettings().model : capability.reqKey;
-  let job = await createGenerationJob({
-    ownerId,
-    externalClientId: input.externalClientId,
-    capability: input.capability,
-    provider: mock ? "mock" : engine === "bailian" ? "bailian" : engine === "evolink" ? "evolink" : "volcengine-visual",
-    reqKey,
-    status: missingBailianKey ? "failed" : "queued",
-    prompt: input.prompt,
-    inputAssetIds: input.inputAssetIds || [],
-    inputUrls: normalizedUrls,
-    outputAssetIds: [],
-    requestPayload: {
-      engine,
-      input,
-      providerPayload
-    },
-    error: missingBailianKey ? { message: "缺少 BAILIAN_API_KEY，请先在设置页配置阿里云百炼 API Key。", retryable: true } : undefined,
-    retryOf: input.retryOf,
-    idempotencyKey: input.idempotencyKey,
-    idempotencyFingerprint: input.idempotencyFingerprint,
-    priority: input.priority,
-    maxAttempts: input.maxAttempts,
-    webhookUrl: input.webhookUrl
-  });
-
-  return job;
+  const provider = mock ? "mock" : engine === "bailian" ? "bailian" : engine === "evolink" ? "evolink" : "volcengine-visual";
+  const requestPayload = {
+    engine,
+    input,
+    providerPayload
+  };
+  return { ownerId, normalizedUrls, engine, providerPayload, missingBailianKey, provider, reqKey, requestPayload };
 }
 
 export async function advanceImageJob(jobId: string, origin: string): Promise<GenerationJob> {
-  const job = await getGenerationJob(jobId);
+  let job = await getGenerationJob(jobId);
   if (!job) throw new Error(`Generation job not found: ${jobId}`);
   if (["succeeded", "failed", "expired", "cancelled"].includes(job.status)) return job;
+  if (job.billing?.status === "pending") job = await chargeGenerationJob(job);
   if (job.provider === "mock") return completeMockJob(job, origin);
   if (!job.providerTaskId) return dispatchImageJob(job);
   return syncImageJob(job.id, origin);
@@ -245,13 +301,7 @@ export async function syncImageJob(jobId: string, origin: string): Promise<Gener
     }));
   }
   if (assets.length) {
-    await recordUsageEvent({
-      ownerId: job.ownerId,
-      jobId: job.id,
-      capability: job.capability,
-      quantity: assets.length,
-      estimatedUnit: "image"
-    });
+    await recordUsageForJob(job);
   }
   return updateGenerationJob(job.id, {
     status: "succeeded",
@@ -276,7 +326,7 @@ async function syncBailianImageJob(job: GenerationJob, origin: string): Promise<
   for (let index = 0; index < urls.length; index += 1) {
     assets.push(await importRemoteImageAsAsset({ ownerId: job.ownerId, url: urls[index], origin, source: sourceForCapability(job.capability), capability: job.capability, jobId: job.id, index, tags: assetTagsForJob(job) }));
   }
-  await recordUsageEvent({ ownerId: job.ownerId, jobId: job.id, capability: job.capability, quantity: assets.length, estimatedUnit: "image" });
+  await recordUsageForJob(job);
   return updateGenerationJob(job.id, { status: "succeeded", outputAssetIds: assets.map((asset) => asset.id), responsePayload: response });
 }
 
@@ -350,13 +400,7 @@ async function syncEvolinkImageJob(job: GenerationJob, origin: string): Promise<
       tags: assetTagsForJob(job)
     }));
   }
-  await recordUsageEvent({
-    ownerId: job.ownerId,
-    jobId: job.id,
-    capability: job.capability,
-    quantity: assets.length,
-    estimatedUnit: "image"
-  });
+  await recordUsageForJob(job);
   return updateGenerationJob(job.id, {
     status: "succeeded",
     outputAssetIds: assets.map((asset) => asset.id),
@@ -364,7 +408,12 @@ async function syncEvolinkImageJob(job: GenerationJob, origin: string): Promise<
   });
 }
 
-export async function retryImageJob(jobId: string, origin: string, ownerId?: string): Promise<GenerationJob> {
+export async function retryImageJob(
+  jobId: string,
+  origin: string,
+  ownerId?: string,
+  usageContext?: UsageContext
+): Promise<GenerationJob> {
   const job = await getGenerationJob(jobId);
   if (!job) throw new Error(`Generation job not found: ${jobId}`);
   if (ownerId && job.ownerId !== ownerId) throw new Error(`Generation job not found: ${jobId}`);
@@ -373,6 +422,7 @@ export async function retryImageJob(jobId: string, origin: string, ownerId?: str
     ...input,
     ownerId: ownerId || job.ownerId,
     capability: job.capability as EnabledImageCapability,
+    usageContext: usageContext || job.usageContext || input.usageContext,
     retryOf: job.id
   }, origin);
 }
@@ -383,8 +433,7 @@ async function completeMockJob(job: GenerationJob, origin: string): Promise<Gene
   const buffer = createMockImageBuffer({
     title: "智念AIGC生成结果",
     prompt: job.prompt || requestInput.prompt,
-    capability: job.capability,
-    resolution: requestInput.resolution
+    capability: job.capability
   });
   const asset = await saveGeneratedAsset({
     ownerId: job.ownerId,
@@ -400,13 +449,6 @@ async function completeMockJob(job: GenerationJob, origin: string): Promise<Gene
       mock: true
     }
   });
-  await recordUsageEvent({
-    ownerId: job.ownerId,
-    jobId: job.id,
-    capability: job.capability,
-    quantity: 1,
-    estimatedUnit: "image"
-  });
   return updateGenerationJob(job.id, {
     status: "succeeded",
     outputAssetIds: [asset.id],
@@ -420,9 +462,7 @@ async function completeMockJob(job: GenerationJob, origin: string): Promise<Gene
   });
 }
 
-function sourceForCapability(capability: string) {
-  if (capability === "image.inpaint") return "edited";
-  if (capability === "image.upscale") return "upscaled";
+function sourceForCapability(_capability: string): "generated" {
   return "generated";
 }