feat(console,schemas): add grant context to custom jwt (#6184)

packages/console/scripts/generate-jwt-customizer-type-definition.ts

@@ -3,6 +3,7 @@ import fs from 'node:fs';
 import {
   accessTokenPayloadGuard,
   clientCredentialsPayloadGuard,
+  jwtCustomizerGrantContextGuard,
   jwtCustomizerUserContextGuard,
 } from '@logto/schemas';
 import prettier from 'prettier';
@@ -13,6 +14,7 @@ const filePath = 'src/consts/jwt-customizer-type-definition.ts';
 
 const typeIdentifiers = `export enum JwtCustomizerTypeDefinitionKey {
   JwtCustomizerUserContext = 'JwtCustomizerUserContext',
+  JwtCustomizerGrantContext = 'JwtCustomizerGrantContext',
   AccessTokenPayload = 'AccessTokenPayload',
   ClientCredentialsPayload = 'ClientCredentialsPayload',
   EnvironmentVariables = 'EnvironmentVariables',
@@ -43,6 +45,11 @@ const createJwtCustomizerTypeDefinitions = async () => {
     'JwtCustomizerUserContext'
   );
 
+  const jwtCustomizerGrantContextTypeDefinition = inferTsDefinitionFromZod(
+    jwtCustomizerGrantContextGuard,
+    'JwtCustomizerGrantContext'
+  );
+
   const accessTokenPayloadTypeDefinition = inferTsDefinitionFromZod(
     accessTokenPayloadGuard,
     'AccessTokenPayload'
@@ -58,6 +65,8 @@ ${typeIdentifiers}
 
 export const jwtCustomizerUserContextTypeDefinition = \`${jwtCustomizerUserContextTypeDefinition}\`;
 
+export const jwtCustomizerGrantContextTypeDefinition = \`${jwtCustomizerGrantContextTypeDefinition}\`;
+
 export const accessTokenPayloadTypeDefinition = \`${accessTokenPayloadTypeDefinition}\`;
 
 export const clientCredentialsPayloadTypeDefinition = \`${clientCredentialsPayloadTypeDefinition}\`;

packages/console/src/pages/CustomizeJwtDetails/MainContent/SettingsSection/InstructionTab/GuideCard/index.tsx

@@ -9,6 +9,7 @@ import * as styles from './index.module.scss';
 
 export enum CardType {
   UserData = 'user_data',
+  GrantData = 'grant_data',
   TokenData = 'token_data',
   FetchExternalData = 'fetch_external_data',
   EnvironmentVariables = 'environment_variables',

packages/console/src/pages/CustomizeJwtDetails/MainContent/SettingsSection/InstructionTab/index.tsx

@@ -16,6 +16,7 @@ import {
   accessTokenPayloadTypeDefinition,
   clientCredentialsPayloadTypeDefinition,
   jwtCustomizerUserContextTypeDefinition,
+  jwtCustomizerGrantContextTypeDefinition,
 } from '@/pages/CustomizeJwtDetails/utils/type-definitions';
 
 import * as tabContentStyles from '../index.module.scss';
@@ -77,6 +78,24 @@ function InstructionTab({ isActive }: Props) {
           />
         </GuideCard>
       )}
+      {tokenType === LogtoJwtTokenKeyType.AccessToken && (
+        <GuideCard
+          name={CardType.GrantData}
+          isExpanded={expendCard === CardType.GrantData}
+          setExpanded={(expand) => {
+            setExpendCard(expand ? CardType.GrantData : undefined);
+          }}
+        >
+          <Editor
+            language="typescript"
+            className={styles.sampleCode}
+            value={jwtCustomizerGrantContextTypeDefinition}
+            height="400px"
+            theme="logto-dark"
+            options={typeDefinitionCodeEditorOptions}
+          />
+        </GuideCard>
+      )}
       <GuideCard
         name={CardType.FetchExternalData}
         isExpanded={expendCard === CardType.FetchExternalData}

packages/console/src/pages/CustomizeJwtDetails/utils/config.tsx

@@ -1,7 +1,9 @@
-import type {
-  AccessTokenPayload,
-  ClientCredentialsPayload,
-  JwtCustomizerUserContext,
+import {
+  GrantType,
+  type AccessTokenPayload,
+  type ClientCredentialsPayload,
+  type JwtCustomizerUserContext,
+  type JwtCustomizerGrantContext,
 } from '@logto/schemas';
 import { type EditorProps } from '@monaco-editor/react';
 
@@ -27,6 +29,7 @@ declare interface CustomJwtClaims extends Record<string, any> {}
  */
 declare type Context = {
   user: ${JwtCustomizerTypeDefinitionKey.JwtCustomizerUserContext};
+  grant?: ${JwtCustomizerTypeDefinitionKey.JwtCustomizerGrantContext};
 }
 
 declare type Payload = {
@@ -199,8 +202,16 @@ const defaultUserContext: Partial<JwtCustomizerUserContext> = {
   organizationRoles: [],
 };
 
+const defaultGrantContext: Partial<JwtCustomizerGrantContext> = {
+  type: GrantType.TokenExchange,
+  subjectTokenContext: {
+    foo: 'bar',
+  },
+};
+
 export const defaultUserTokenContextData = {
   user: defaultUserContext,
+  grant: defaultGrantContext,
 };
 
 export const accessTokenPayloadTestModel: ModelSettings = {
@@ -223,6 +234,6 @@ export const userContextTestModel: ModelSettings = {
   language: 'json',
   icon: <UserFileIcon />,
   name: 'user-token-context.json',
-  title: 'User data',
+  title: 'Context data',
   defaultValue: JSON.stringify(defaultUserTokenContextData, null, 2),
 };

packages/console/src/pages/CustomizeJwtDetails/utils/type-definitions.ts

@@ -3,6 +3,7 @@ import {
   accessTokenPayloadTypeDefinition,
   clientCredentialsPayloadTypeDefinition,
   jwtCustomizerUserContextTypeDefinition,
+  jwtCustomizerGrantContextTypeDefinition,
 } from '@/consts/jwt-customizer-type-definition';
 
 import { type JwtCustomizerForm } from '../type';
@@ -12,11 +13,14 @@ export {
   accessTokenPayloadTypeDefinition,
   clientCredentialsPayloadTypeDefinition,
   jwtCustomizerUserContextTypeDefinition,
+  jwtCustomizerGrantContextTypeDefinition,
 } from '@/consts/jwt-customizer-type-definition';
 
 export const buildAccessTokenJwtCustomizerContextTsDefinition = () => {
   return `declare ${jwtCustomizerUserContextTypeDefinition}
 
+  declare ${jwtCustomizerGrantContextTypeDefinition}
+
   declare ${accessTokenPayloadTypeDefinition}`;
 };
 

packages/phrases/src/locales/en/translation/admin-console/jwt-claims.ts

@@ -30,6 +30,11 @@ const jwt_claims = {
     title: 'User data',
     subtitle: 'Use `data.user` input parameter to provide vital user info.',
   },
+  grant_data: {
+    title: 'Grant data',
+    subtitle:
+      'Use `data.grant` input parameter to provide vital grant info, only available for token exchange.',
+  },
   token_data: {
     title: 'Token data',
     subtitle: 'Use `token` input parameter for current access token payload. ',

packages/schemas/src/types/logto-config/jwt-customizer.ts

@@ -10,6 +10,7 @@ import {
   type UserSsoIdentity,
 } from '../../db-entries/index.js';
 import { mfaFactorsGuard, type MfaFactors } from '../../foundations/index.js';
+import { GrantType } from '../oidc-config.js';
 import { scopeResponseGuard, type ScopeResponse } from '../scope.js';
 import { userInfoGuard, type UserInfo } from '../user.js';
 
@@ -67,11 +68,23 @@ export const jwtCustomizerUserContextGuard = userInfoGuard.extend({
     .array(),
 }) satisfies ZodType<JwtCustomizerUserContext>;
 
+export const jwtCustomizerGrantContextGuard = z.object({
+  type: z.literal(GrantType.TokenExchange), // Only support token exchange for now
+  subjectTokenContext: jsonObjectGuard,
+});
+
+export type JwtCustomizerGrantContext = z.infer<typeof jwtCustomizerGrantContextGuard>;
+
 export const accessTokenJwtCustomizerGuard = jwtCustomizerGuard
   .extend({
     // Use partial token guard since users customization may not rely on all fields.
     tokenSample: accessTokenPayloadGuard.partial().optional(),
-    contextSample: z.object({ user: jwtCustomizerUserContextGuard.partial() }).optional(),
+    contextSample: z
+      .object({
+        user: jwtCustomizerUserContextGuard.partial(),
+        grant: jwtCustomizerGrantContextGuard.partial().optional(),
+      })
+      .optional(),
   })
   .strict();