feat(server): refactor for byok (#14911)

packages/backend/server/src/models/copilot-byok.ts

@@ -0,0 +1,139 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+
+import { BaseModel } from './base';
+
+export type UpsertAiWorkspaceByokConfigInput = {
+  id?: string | null;
+  workspaceId: string;
+  provider: string;
+  name: string;
+  description: string | null;
+  encryptedApiKey?: string;
+  endpoint: string | null;
+  sortOrder: number;
+  enabled: boolean;
+  userId?: string;
+};
+
+@Injectable()
+export class CopilotWorkspaceByokConfigModel extends BaseModel {
+  async list(workspaceId: string) {
+    return await this.db.aiWorkspaceByokConfig.findMany({
+      where: { workspaceId },
+      orderBy: [{ sortOrder: 'asc' }, { createdAt: 'asc' }],
+    });
+  }
+
+  async listEnabled(workspaceId: string) {
+    return await this.db.aiWorkspaceByokConfig.findMany({
+      where: { workspaceId, enabled: true },
+      orderBy: [{ sortOrder: 'asc' }, { createdAt: 'asc' }],
+    });
+  }
+
+  async get(id: string) {
+    return await this.db.aiWorkspaceByokConfig.findUnique({
+      where: { id },
+    });
+  }
+
+  @Transactional()
+  async upsert(input: UpsertAiWorkspaceByokConfigInput) {
+    const data = {
+      provider: input.provider,
+      name: input.name,
+      description: input.description,
+      endpoint: input.endpoint,
+      sortOrder: input.sortOrder,
+      enabled: input.enabled,
+      updatedBy: input.userId,
+      ...(input.encryptedApiKey
+        ? {
+            encryptedApiKey: input.encryptedApiKey,
+            lastValidatedAt: new Date(),
+            lastValidationError: null,
+            disabledReason: null,
+            lastError: null,
+            lastErrorAt: null,
+          }
+        : {}),
+    };
+
+    return input.id
+      ? await this.db.aiWorkspaceByokConfig.update({
+          where: { id: input.id, workspaceId: input.workspaceId },
+          data,
+        })
+      : await this.db.aiWorkspaceByokConfig.create({
+          data: {
+            ...data,
+            encryptedApiKey: input.encryptedApiKey ?? '',
+            workspaceId: input.workspaceId,
+            createdBy: input.userId,
+          },
+        });
+  }
+
+  @Transactional()
+  async reorder(workspaceId: string, ids: string[], userId?: string) {
+    await Promise.all(
+      ids.map((id, sortOrder) =>
+        this.db.aiWorkspaceByokConfig.update({
+          where: { id, workspaceId },
+          data: { sortOrder, updatedBy: userId },
+        })
+      )
+    );
+  }
+
+  @Transactional()
+  async delete(workspaceId: string, id: string) {
+    await this.db.aiWorkspaceByokConfig.delete({ where: { id, workspaceId } });
+  }
+
+  @Transactional()
+  async clear(workspaceId: string, provider?: string | null) {
+    await this.db.aiWorkspaceByokConfig.deleteMany({
+      where: { workspaceId, ...(provider ? { provider } : {}) },
+    });
+  }
+
+  @Transactional()
+  async markValidated(workspaceId: string, id: string, userId?: string) {
+    await this.db.aiWorkspaceByokConfig.update({
+      where: { id, workspaceId },
+      data: {
+        enabled: true,
+        disabledReason: null,
+        lastValidatedAt: new Date(),
+        lastValidationError: null,
+        lastError: null,
+        lastErrorAt: null,
+        updatedBy: userId,
+      },
+    });
+  }
+
+  @Transactional()
+  async markFailure(workspaceId: string, id: string, message: string) {
+    await this.db.aiWorkspaceByokConfig.update({
+      where: { id, workspaceId },
+      data: {
+        enabled: false,
+        disabledReason: 'recent_failure',
+        lastValidationError: message,
+        lastError: message,
+        lastErrorAt: new Date(),
+      },
+    });
+  }
+
+  @Transactional()
+  async touchUsed(workspaceId: string, id: string) {
+    await this.db.aiWorkspaceByokConfig.updateMany({
+      where: { id, workspaceId },
+      data: { lastUsedAt: new Date() },
+    });
+  }
+}

packages/backend/server/src/models/copilot-session.ts

@@ -1005,20 +1005,26 @@ export class CopilotSessionModel extends BaseModel {
       .filter(({ promptAction }) => !promptAction)
       .map(({ messageCost }) => messageCost)
       .reduce((prev, cost) => prev + cost, 0);
-    const [actionRunCost, legacyActionSessionCost, transcriptSettlementCost] =
-      await Promise.all([
-        this.models.copilotActionRun.countSucceededByUser(userId),
-        this.models.copilotActionRun.countLegacyPromptActionSessionsWithoutRun(
-          userId
-        ),
-        this.models.copilotTranscriptTask.countSettledByUser(userId),
-      ]);
-    return (
+    const [
+      actionRunCost,
+      legacyActionSessionCost,
+      transcriptSettlementCost,
+      byokQuotaExemptCost,
+    ] = await Promise.all([
+      this.models.copilotActionRun.countSucceededByUser(userId),
+      this.models.copilotActionRun.countLegacyPromptActionSessionsWithoutRun(
+        userId
+      ),
+      this.models.copilotTranscriptTask.countSettledByUser(userId),
+      this.models.copilotUsage.countQuotaExemptByokUsage(userId),
+    ]);
+    const quotaBackedCost =
       regularMessageCost +
       actionRunCost +
       legacyActionSessionCost +
-      transcriptSettlementCost
-    );
+      transcriptSettlementCost -
+      byokQuotaExemptCost;
+    return Math.max(0, quotaBackedCost);
   }
 
   async cleanupEmptySessions(earlyThen: Date) {

packages/backend/server/src/models/copilot-usage.ts

@@ -0,0 +1,142 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { Prisma } from '@prisma/client';
+
+import { BaseModel } from './base';
+
+type CreateAiUsageEventInput = {
+  workspaceId: string;
+  userId?: string;
+  provider: string;
+  providerSource: string;
+  featureKind: string;
+  model?: string | null;
+  sessionId?: string;
+  taskId?: string;
+  actionId?: string;
+  billingUnitId?: string;
+  promptTokens?: number;
+  completionTokens?: number;
+  totalTokens?: number;
+  cachedTokens?: number;
+};
+
+type UsageAggregateRow = {
+  date: string;
+  featureKind: string;
+  totalTokens: number | bigint | null;
+};
+
+type CountRow = {
+  count: number | bigint;
+};
+
+const BYOK_PROVIDER_SOURCES = ['byok_server', 'byok_local'];
+const QUOTA_EXEMPT_BYOK_FEATURES = ['chat', 'action', 'image', 'transcript'];
+
+@Injectable()
+export class CopilotUsageModel extends BaseModel {
+  @Transactional()
+  async create(input: CreateAiUsageEventInput) {
+    await this.db.aiUsageEvent.create({
+      data: {
+        workspaceId: input.workspaceId,
+        userId: input.userId,
+        provider: input.provider,
+        providerSource: input.providerSource,
+        featureKind: input.featureKind,
+        model: input.model ?? null,
+        sessionId: input.sessionId,
+        taskId: input.taskId,
+        actionId: input.actionId,
+        billingUnitId: input.billingUnitId,
+        promptTokens: input.promptTokens ?? 0,
+        completionTokens: input.completionTokens ?? 0,
+        totalTokens: input.totalTokens ?? 0,
+        cachedTokens: input.cachedTokens ?? 0,
+      },
+    });
+  }
+
+  async countQuotaExemptByokUsage(userId: string) {
+    const rows = await this.db.$queryRaw<CountRow[]>(Prisma.sql`
+      WITH "byok_usage" AS (
+        SELECT "billing_unit_id", "feature_kind"
+        FROM "ai_usage_events"
+        WHERE "user_id" = ${userId}
+          AND "provider_source" IN (${Prisma.join(BYOK_PROVIDER_SOURCES)})
+          AND "feature_kind" IN (${Prisma.join(QUOTA_EXEMPT_BYOK_FEATURES)})
+          AND "billing_unit_id" IS NOT NULL
+      ),
+      "message_units" AS (
+        SELECT DISTINCT "usage"."billing_unit_id"
+        FROM "byok_usage" AS "usage"
+        JOIN "ai_sessions_messages" AS "message"
+          ON "message"."id" = "usage"."billing_unit_id"
+        JOIN "ai_sessions_metadata" AS "session"
+          ON "session"."id" = "message"."session_id"
+        WHERE "usage"."feature_kind" IN ('chat', 'action', 'image')
+          AND "message"."role" = 'user'
+          AND "session"."user_id" = ${userId}
+          AND ("session"."prompt_action" IS NULL OR "session"."prompt_action" = '')
+      ),
+      "action_units" AS (
+        SELECT DISTINCT "usage"."billing_unit_id"
+        FROM "byok_usage" AS "usage"
+        JOIN "ai_action_runs" AS "run"
+          ON "run"."id" = "usage"."billing_unit_id"
+        WHERE "usage"."feature_kind" IN ('action', 'image')
+          AND "run"."user_id" = ${userId}
+          AND "run"."status" = 'succeeded'
+          AND "run"."action_id" NOT LIKE 'transcript.audio.%'
+      ),
+      "transcript_units" AS (
+        SELECT DISTINCT "usage"."billing_unit_id"
+        FROM "byok_usage" AS "usage"
+        JOIN "ai_transcript_tasks" AS "task"
+          ON "task"."id" = "usage"."billing_unit_id"
+        WHERE "usage"."feature_kind" = 'transcript'
+          AND "task"."user_id" = ${userId}
+          AND "task"."status" = 'settled'
+      )
+      SELECT (
+        (SELECT COUNT(*) FROM "message_units") +
+        (SELECT COUNT(*) FROM "action_units") +
+        (SELECT COUNT(*) FROM "transcript_units")
+      ) AS "count"
+    `);
+    const count = rows[0]?.count ?? 0;
+    return typeof count === 'bigint' ? Number(count) : count;
+  }
+
+  async aggregateByDay(input: {
+    workspaceId: string;
+    from: Date;
+    to: Date;
+    providerSources: string[];
+  }) {
+    if (!input.providerSources.length) return [];
+
+    const rows = await this.db.$queryRaw<UsageAggregateRow[]>(Prisma.sql`
+      SELECT
+        to_char(date_trunc('day', "created_at" AT TIME ZONE 'UTC'), 'YYYY-MM-DD') AS "date",
+        "feature_kind" AS "featureKind",
+        COALESCE(SUM("total_tokens"), 0)::bigint AS "totalTokens"
+      FROM "ai_usage_events"
+      WHERE "workspace_id" = ${input.workspaceId}
+        AND "provider_source" IN (${Prisma.join(input.providerSources)})
+        AND "created_at" >= ${input.from}
+        AND "created_at" < ${input.to}
+      GROUP BY 1, 2
+      ORDER BY 1 ASC, 2 ASC
+    `);
+
+    return rows.map(row => {
+      return {
+        date: new Date(`${row.date}T00:00:00.000Z`),
+        featureKind: row.featureKind,
+        totalTokens: Number(row.totalTokens ?? 0),
+      };
+    });
+  }
+}

packages/backend/server/src/models/index.ts

@@ -17,10 +17,12 @@ import { CommentModel } from './comment';
 import { CommentAttachmentModel } from './comment-attachment';
 import { AppConfigModel } from './config';
 import { CopilotActionRunModel } from './copilot-action-run';
+import { CopilotWorkspaceByokConfigModel } from './copilot-byok';
 import { CopilotContextModel } from './copilot-context';
 import { CopilotJobModel } from './copilot-job';
 import { CopilotSessionModel } from './copilot-session';
 import { CopilotTranscriptTaskModel } from './copilot-transcript-task';
+import { CopilotUsageModel } from './copilot-usage';
 import { CopilotWorkspaceConfigModel } from './copilot-workspace';
 import { DocModel } from './doc';
 import { DocUserModel } from './doc-user';
@@ -58,10 +60,12 @@ const MODELS = {
   notification: NotificationModel,
   userSettings: UserSettingsModel,
   copilotSession: CopilotSessionModel,
+  copilotUsage: CopilotUsageModel,
   copilotTranscriptTask: CopilotTranscriptTaskModel,
   copilotActionRun: CopilotActionRunModel,
   copilotContext: CopilotContextModel,
   copilotWorkspace: CopilotWorkspaceConfigModel,
+  copilotWorkspaceByokConfig: CopilotWorkspaceByokConfigModel,
   copilotJob: CopilotJobModel,
   appConfig: AppConfigModel,
   comment: CommentModel,
@@ -133,10 +137,12 @@ export * from './calendar-subscription';
 export * from './comment';
 export * from './comment-attachment';
 export * from './common';
+export * from './copilot-byok';
 export * from './copilot-context';
 export * from './copilot-job';
 export * from './copilot-session';
 export * from './copilot-transcript-task';
+export * from './copilot-usage';
 export * from './copilot-workspace';
 export * from './doc';
 export * from './doc-user';