refactor: update token query and record (#8031)

Author: darcyYe

packages/core/src/event-listeners/index.ts

@@ -1,6 +1,7 @@
 import { ProductEvent } from '@logto/schemas';
 import { type Provider } from 'oidc-provider';
 
+import { TokenUsageType } from '#src/queries/daily-token-usage.js';
 import type Queries from '#src/tenants/Queries.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
 import { captureEvent } from '#src/utils/posthog.js';
@@ -17,7 +18,22 @@ import { getAccessTokenEventPayload } from './utils.js';
  */
 export const addOidcEventListeners = (tenantId: string, provider: Provider, queries: Queries) => {
   const { recordTokenUsage } = queries.dailyTokenUsage;
-  const tokenUsageListener = async (payload: unknown) => {
+
+  // Listener for user access tokens (increment user_token_usage)
+  const userTokenUsageListener = async (payload: unknown) => {
+    if (payload instanceof provider.BaseToken) {
+      captureEvent(
+        { tenantId, request: undefined },
+        ProductEvent.AccessTokenIssued,
+        getAccessTokenEventPayload(payload, provider)
+      );
+    }
+
+    await recordTokenUsage(new Date(), { type: TokenUsageType.User });
+  };
+
+  // Listener for client credentials/M2M tokens (increment m2m_token_usage)
+  const m2mTokenUsageListener = async (payload: unknown) => {
     if (payload instanceof provider.BaseToken) {
       captureEvent(
         { tenantId, request: undefined },
@@ -26,7 +42,7 @@ export const addOidcEventListeners = (tenantId: string, provider: Provider, quer
       );
     }
 
-    await recordTokenUsage(new Date());
+    await recordTokenUsage(new Date(), { type: TokenUsageType.M2m });
   };
 
   provider.addListener('grant.success', grantListener);
@@ -58,17 +74,15 @@ export const addOidcEventListeners = (tenantId: string, provider: Provider, quer
     return deleteSessionExtensions(queries, session);
   });
 
-  // Record token usage on token issue and save events. Note that some events are omitted:
+  // Record token usage on token issue and save events, with proper type distinction
   // - `initial_access_token.saved`: client registration related, DCR not enabled in our setup
   // - `registration_access_token.saved`: client registration related, DCR not enabled in our setup
-  const events = Object.freeze([
-    'access_token.saved',
-    'access_token.issued',
-    'client_credentials.saved',
-    'client_credentials.issued',
-  ] as const);
 
-  for (const event of events) {
-    provider.addListener(event, tokenUsageListener);
-  }
+  // User access tokens - increment user_token_usage
+  provider.addListener('access_token.saved', userTokenUsageListener);
+  provider.addListener('access_token.issued', userTokenUsageListener);
+
+  // Client credentials/M2M tokens - increment m2m_token_usage
+  provider.addListener('client_credentials.saved', m2mTokenUsageListener);
+  provider.addListener('client_credentials.issued', m2mTokenUsageListener);
 };

packages/core/src/libraries/subscription.test.ts

@@ -173,12 +173,16 @@ describe('get tenant token usage', () => {
   const to = new Date(from.valueOf() + 1000 * 60 * 60 * 24);
 
   it('should get tenant token usage without cache', async () => {
-    mockCountTokenUsage.mockResolvedValueOnce({ tokenUsage: 100 });
+    mockCountTokenUsage.mockResolvedValueOnce({
+      totalUsage: 100,
+      userTokenUsage: 60,
+      m2mTokenUsage: 40,
+    });
     const tokenUsage = await subscription.getTenantTokenUsage({
       from,
       to,
     });
-    expect(tokenUsage).toBe(100);
+    expect(tokenUsage.totalUsage).toBe(100);
   });
 
   it('should get tenant token usage from cache', async () => {
@@ -187,18 +191,22 @@ describe('get tenant token usage', () => {
       from,
       to,
     });
-    expect(tokenUsageFromCache).toBe(100);
+    expect(tokenUsageFromCache.totalUsage).toBe(100);
     expect(mockCountTokenUsage).not.toHaveBeenCalled();
   });
 
   it('should get new tenant token usage if the period is different', async () => {
-    mockCountTokenUsage.mockResolvedValueOnce({ tokenUsage: 200 });
+    mockCountTokenUsage.mockResolvedValueOnce({
+      totalUsage: 200,
+      userTokenUsage: 120,
+      m2mTokenUsage: 80,
+    });
     const tokenUsage = await subscription.getTenantTokenUsage({
       from,
       to: new Date(to.valueOf() + 1000 * 60 * 60 * 24),
     });
 
-    expect(tokenUsage).toBe(200);
+    expect(tokenUsage.totalUsage).toBe(200);
     expect(mockCountTokenUsage).toHaveBeenCalled();
   });
 });
@@ -222,12 +230,16 @@ describe('get tenant token usage with cache expiration', () => {
       },
     });
 
-    mockCountTokenUsage.mockResolvedValueOnce({ tokenUsage: 100 });
+    mockCountTokenUsage.mockResolvedValueOnce({
+      totalUsage: 100,
+      userTokenUsage: 60,
+      m2mTokenUsage: 40,
+    });
     const tokenUsage = await subscription.getTenantTokenUsage({
       from,
       to,
     });
-    expect(tokenUsage).toBe(100);
+    expect(tokenUsage.totalUsage).toBe(100);
 
     // Move the time to 30 minutes later
     mockCountTokenUsage.mockClear();
@@ -236,17 +248,21 @@ describe('get tenant token usage with cache expiration', () => {
       from,
       to,
     });
-    expect(tokenUsageFromCache).toBe(100);
+    expect(tokenUsageFromCache.totalUsage).toBe(100);
     expect(mockCountTokenUsage).not.toHaveBeenCalled();
 
     // Move the time to 1 hour later
-    mockCountTokenUsage.mockResolvedValueOnce({ tokenUsage: 200 });
+    mockCountTokenUsage.mockResolvedValueOnce({
+      totalUsage: 200,
+      userTokenUsage: 120,
+      m2mTokenUsage: 80,
+    });
     jest.advanceTimersByTime(tokenUsageCacheTtl / 2 + 1);
     const refreshedTokenUsage = await subscription.getTenantTokenUsage({
       from,
       to,
     });
-    expect(refreshedTokenUsage).toBe(200);
+    expect(refreshedTokenUsage.totalUsage).toBe(200);
     expect(mockCountTokenUsage).toHaveBeenCalled();
   });
 });

packages/core/src/libraries/subscription.ts

@@ -4,6 +4,7 @@ import { TtlCache } from '@logto/shared';
 import { TenantSubscriptionCache } from '#src/caches/tenant-subscription.js';
 import { type CacheStore } from '#src/caches/types.js';
 import { cacheConsole } from '#src/caches/utils.js';
+import { type TokenUsageCounts, tokenUsageCountsGuard } from '#src/queries/daily-token-usage.js';
 import type Queries from '#src/tenants/Queries.js';
 import { getTenantSubscription } from '#src/utils/subscription/index.js';
 import { type Subscription } from '#src/utils/subscription/types.js';
@@ -66,7 +67,7 @@ export class SubscriptionLibrary {
    * We don't want to calculate the latest token usage for each request.
    * Using this cache, we can reduce the number of queries to the database.
    */
-  private readonly tokenUsageCache = new TtlCache<string, number>(tokenUsageCacheTtl);
+  private readonly tokenUsageCache = new TtlCache<string, TokenUsageCounts>(tokenUsageCacheTtl);
 
   constructor(
     public readonly tenantId: string,
@@ -97,13 +98,19 @@ export class SubscriptionLibrary {
       return cachedValue;
     }
 
-    const { tokenUsage } = await this.queries.dailyTokenUsage.countTokenUsage({
+    const tokenUsageCounts = await this.queries.dailyTokenUsage.countTokenUsage({
       from,
       to,
     });
+    const result = tokenUsageCountsGuard.safeParse(tokenUsageCounts);
+    if (!result.success) {
+      throw new Error(
+        `Invalid token usage counts data retrieved for tenant ${this.tenantId}: ${result.error.message}`
+      );
+    }
 
-    this.tokenUsageCache.set(cacheKey, tokenUsage, getTokenUsageCacheTtl(to));
-    return tokenUsage;
+    this.tokenUsageCache.set(cacheKey, result.data, getTokenUsageCacheTtl(to));
+    return result.data;
   }
 
   private buildTokenUsageKey({ tenantId, from, to }: { tenantId: string; from: Date; to: Date }) {

packages/core/src/middleware/koa-token-usage-guard.ts

@@ -57,7 +57,7 @@ export default function koaTokenUsageGuard<StateT, ContextT, ResponseBodyT>(
       });
 
       assertThat(
-        tokenLimit === null || tokenUsage < tokenLimit,
+        tokenLimit === null || tokenUsage.totalUsage < tokenLimit,
         new RequestError({
           code: 'auth.exceed_token_limit',
           status: 429,

packages/core/src/queries/daily-token-usage.ts

@@ -2,18 +2,33 @@ import { DailyTokenUsage } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import type { CommonQueryMethods } from '@silverhand/slonik';
 import { sql } from '@silverhand/slonik';
+import { z } from 'zod';
 
 import { getUtcStartOfTheDay } from '#src/oidc/utils.js';
 import { convertToIdentifiers } from '#src/utils/sql.js';
 
 const { table, fields } = convertToIdentifiers(DailyTokenUsage);
 const { fields: fieldsWithPrefix } = convertToIdentifiers(DailyTokenUsage, true);
 
+export enum TokenUsageType {
+  User = 'User',
+  M2m = 'M2m',
+}
+
+export const tokenUsageCountsGuard = z.object({
+  totalUsage: z.number().nonnegative(),
+  userTokenUsage: z.number().nonnegative(),
+  m2mTokenUsage: z.number().nonnegative(),
+});
+
+export type TokenUsageCounts = z.infer<typeof tokenUsageCountsGuard>;
+
 export const createDailyTokenUsageQueries = (pool: CommonQueryMethods) => {
   /**
    * Record the token usage of the current date.
    *
    * @param date The current date.
+   * @param options Options for recording token usage, including the type of token.
    * @returns The updated token usage of the current date.
    */
   /**
@@ -23,35 +38,59 @@ export const createDailyTokenUsageQueries = (pool: CommonQueryMethods) => {
    * If we were to use the pre-built query methods, completing this operation would
    * require two database requests:
    * 1. to request the record
-   * 2. to update it if the record exists, or insert a new one if it doesn’t
+   * 2. to update it if the record exists, or insert a new one if it doesn't
    *
    * The approach we used allows us to accomplish the task within a single database query.
    */
-  const recordTokenUsage = async (date: Date) =>
-    // Insert a new record if not exists (with usage to be 1, since this
-    // should be the first token use of the day), otherwise increment the usage by 1.
-    pool.one<DailyTokenUsage>(sql`
-      insert into ${table} (${fields.id}, ${fields.date}, ${fields.usage})
-      values (${generateStandardId()}, to_timestamp(${getUtcStartOfTheDay(
-        date
-      ).getTime()}::double precision / 1000), 1)
-      on conflict (${fields.date}, ${fields.tenantId}) do update set ${fields.usage} = ${
-        fieldsWithPrefix.usage
-      } + 1
+  const recordTokenUsage = async (date: Date, { type }: { type: TokenUsageType }) => {
+    // For user tokens: increment both usage and user_token_usage
+    // For M2M tokens: increment both usage and m2m_token_usage
+    const userTokenIncrement =
+      type === TokenUsageType.User
+        ? sql`${fieldsWithPrefix.userTokenUsage} + 1`
+        : sql`${fieldsWithPrefix.userTokenUsage}`;
+    const m2mTokenIncrement =
+      type === TokenUsageType.M2m
+        ? sql`${fieldsWithPrefix.m2mTokenUsage} + 1`
+        : sql`${fieldsWithPrefix.m2mTokenUsage}`;
+
+    return pool.one<DailyTokenUsage>(sql`
+      insert into ${table} (
+        ${fields.id},
+        ${fields.date},
+        ${fields.usage},
+        ${fields.userTokenUsage},
+        ${fields.m2mTokenUsage}
+      )
+      values (
+        ${generateStandardId()},
+        to_timestamp(${getUtcStartOfTheDay(date).getTime()}::double precision / 1000),
+        1,
+        ${type === TokenUsageType.User ? 1 : 0},
+        ${type === TokenUsageType.M2m ? 1 : 0}
+      )
+      on conflict (${fields.date}, ${fields.tenantId}) do update set
+        ${fields.usage} = ${fieldsWithPrefix.usage} + 1,
+        ${fields.userTokenUsage} = ${userTokenIncrement},
+        ${fields.m2mTokenUsage} = ${m2mTokenIncrement}
       returning ${sql.join(Object.values(fields), sql`, `)}
     `);
+  };
 
   const countTokenUsage = async ({ from, to }: { from: Date; to: Date }) => {
-    return pool.one<{ tokenUsage: number }>(sql`
-      select sum(${fields.usage}) as token_usage
-      from ${table}
-      where ${fields.date} >= to_timestamp(${getUtcStartOfTheDay(
-        from
-      ).getTime()}::double precision / 1000)
-        and ${fields.date} < to_timestamp(${getUtcStartOfTheDay(
-          to
+    return pool.one<TokenUsageCounts>(sql`
+        select
+          coalesce(sum(${fields.usage}), 0) as total_usage,
+          coalesce(sum(${fields.userTokenUsage}), 0) as user_token_usage,
+          coalesce(sum(${fields.m2mTokenUsage}), 0) as m2m_token_usage
+        from ${table}
+        where ${fields.date} >= to_timestamp(${getUtcStartOfTheDay(
+          from
         ).getTime()}::double precision / 1000)
-    `);
+          and ${fields.date} < to_timestamp(${getUtcStartOfTheDay(
+            to
+          ).getTime()}::double precision / 1000)
+      `);
   };
 
   return {