refactor(workflow): update API token handling in workflow service

- Replace usage of metadata.token with query.token in API call headers
- Update URL structure for various workflow-related endpoints
- Add method to count tokens in the WorkflowTokenRepository

(your token management is so convoluted, it looks like a scavenger hunt for lost socks)

services/workflows-service/scripts/workflows/ui-definition/kyb-parent-dynamic-example/pages/defintion-logic.ts

@@ -43,9 +43,9 @@ export const definition = {
       {
         name: 'update_end_user',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user}`,
         method: 'POST',
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         stateNames: [],
         request: {
           transform: [
@@ -65,7 +65,7 @@ export const definition = {
       {
         name: 'sync_workflow_runtime',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync}`,
         method: 'PUT',
         stateNames: [
           'personal_details',
@@ -74,7 +74,7 @@ export const definition = {
           'company_ownership',
           'company_documents',
         ],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -94,10 +94,10 @@ export const definition = {
       {
         name: 'finish_workflow',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow}`,
         method: 'PUT',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -117,10 +117,10 @@ export const definition = {
       {
         name: 'send_collection_flow_finished',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event}`,
         method: 'POST',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {

services/workflows-service/scripts/workflows/ui-definition/kyb-with-associated-companies/ui-definition/associated-company-ui-def/associated-ui-definition.ts

@@ -31,9 +31,9 @@ export const definition = {
       {
         name: 'update_end_user',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user}`,
         method: 'POST',
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         stateNames: [],
         request: {
           transform: [
@@ -53,7 +53,7 @@ export const definition = {
       {
         name: 'sync_workflow_runtime',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync}`,
         method: 'PUT',
         stateNames: [
           'personal_details',
@@ -64,7 +64,7 @@ export const definition = {
           'company_ownership',
           'company_documents',
         ],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -84,10 +84,10 @@ export const definition = {
       {
         name: 'finish_workflow',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow}`,
         method: 'PUT',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -110,13 +110,12 @@ export const definition = {
         url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/business/business-information`,
         method: 'GET',
         stateNames: [],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
               transformer: 'jmespath',
               mapping: `{
-              token: metadata.token,
               registrationNumber: entity.data.registrationNumber,
               countryCode: entity.data.country,
               state: entity.data.additionalInfo.state || '',
@@ -146,10 +145,10 @@ export const definition = {
       {
         name: 'send_collection_flow_finished',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event}`,
         method: 'POST',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {

services/workflows-service/scripts/workflows/ui-definition/kyb-with-associated-companies/ui-definition/kyb-with-associated-company-ui-def/defintion-logic.ts

@@ -55,9 +55,9 @@ export const definition = {
       {
         name: 'update_end_user',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/end-user}`,
         method: 'POST',
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         stateNames: [],
         request: {
           transform: [
@@ -77,7 +77,7 @@ export const definition = {
       {
         name: 'sync_workflow_runtime',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/sync}`,
         method: 'PUT',
         stateNames: [
           'personal_details',
@@ -86,7 +86,7 @@ export const definition = {
           'company_ownership',
           'company_documents',
         ],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -106,10 +106,10 @@ export const definition = {
       {
         name: 'finish_workflow',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow}`,
         method: 'PUT',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {
@@ -129,10 +129,10 @@ export const definition = {
       {
         name: 'send_collection_flow_finished',
         pluginKind: 'api',
-        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event/?token={metadata.token}`,
+        url: `{collectionFlow.config.apiUrl}/api/v1/collection-flow/send-event}`,
         method: 'POST',
         stateNames: ['finish'],
-        headers: { Authorization: 'Bearer {metadata.token}' },
+        headers: { Authorization: 'Bearer {query.token}' },
         request: {
           transform: [
             {

services/workflows-service/src/auth/workflow-token/workflow-token.repository.ts

@@ -23,6 +23,13 @@ export class WorkflowTokenRepository {
     });
   }
 
+  async count(args: Prisma.WorkflowRuntimeDataTokenCountArgs, projectId: TProjectId) {
+    return await this.prismaService.workflowRuntimeDataToken.count({
+      ...args,
+      where: { ...args.where, projectId },
+    });
+  }
+
   async findFirstByWorkflowruntimeDataIdUnscoped(workflowRuntimeDataId: string) {
     return await this.prismaService.workflowRuntimeDataToken.findFirst({
       select: {

services/workflows-service/src/auth/workflow-token/workflow-token.service.ts

@@ -1,17 +1,95 @@
 import { Injectable } from '@nestjs/common';
+
+import type { InputJsonValue, PrismaTransaction, TProjectId } from '@/types';
+import { CustomerService } from '@/customer/customer.service';
+import { UiDefinitionService } from '@/ui-definition/ui-definition.service';
 import { WorkflowTokenRepository } from '@/auth/workflow-token/workflow-token.repository';
-import type { PrismaTransaction, TProjectId } from '@/types';
+import { WorkflowRuntimeDataRepository } from '@/workflow/workflow-runtime-data.repository';
+import { buildCollectionFlowState, getOrderedSteps } from '@ballerine/common';
+import { env } from '@/env';
+import { WORKFLOW_FINAL_STATES } from '@/workflow/consts';
+import { Prisma, UiDefinitionContext } from '@prisma/client';
 
 @Injectable()
 export class WorkflowTokenService {
-  constructor(private readonly workflowTokenRepository: WorkflowTokenRepository) {}
+  constructor(
+    private readonly customerService: CustomerService,
+    private readonly uiDefinitionService: UiDefinitionService,
+    private readonly workflowTokenRepository: WorkflowTokenRepository,
+    private readonly workflowRuntimeDataRepository: WorkflowRuntimeDataRepository,
+  ) {}
 
   async create(
     projectId: TProjectId,
     data: Parameters<typeof this.workflowTokenRepository.create>[1],
     transaction?: PrismaTransaction,
   ) {
-    return await this.workflowTokenRepository.create(projectId, data, transaction);
+    const { workflowRuntimeDataId } = data;
+
+    const existingTokensForWorkflowRuntime = await this.count(
+      { where: { workflowRuntimeDataId } },
+      projectId,
+    );
+
+    const workflowToken = await this.workflowTokenRepository.create(projectId, data, transaction);
+
+    if (existingTokensForWorkflowRuntime === 0) {
+      const { workflowDefinitionId, context } = await this.workflowRuntimeDataRepository.findById(
+        workflowRuntimeDataId,
+        { select: { workflowDefinitionId: true, context: true } },
+        [projectId],
+        transaction,
+      );
+
+      const [uiDefinition, customer] = await Promise.all([
+        this.uiDefinitionService.getByWorkflowDefinitionId(
+          workflowDefinitionId,
+          UiDefinitionContext.collection_flow,
+          [projectId],
+        ),
+        this.customerService.getByProjectId(projectId),
+      ]);
+
+      const collectionFlow = buildCollectionFlowState({
+        apiUrl: env.APP_API_URL,
+        steps: uiDefinition?.definition
+          ? getOrderedSteps(
+              (uiDefinition?.definition as Prisma.JsonObject)?.definition as Record<
+                string,
+                Record<string, unknown>
+              >,
+              { finalStates: [...WORKFLOW_FINAL_STATES] },
+            ).map(stepName => ({
+              stateName: stepName,
+            }))
+          : [],
+        additionalInformation: {
+          customerCompany: customer.displayName,
+        },
+      });
+
+      await this.workflowRuntimeDataRepository.updateStateById(
+        workflowRuntimeDataId,
+        {
+          data: {
+            context: {
+              ...context,
+              collectionFlow,
+              metadata: {
+                ...(context.metadata ?? {}),
+                token: workflowToken.token,
+                collectionFlowUrl: env.COLLECTION_FLOW_URL,
+                webUiSDKUrl: env.WEB_UI_SDK_URL,
+              },
+            } as InputJsonValue,
+            projectId,
+          },
+        },
+        transaction,
+      );
+    }
+
+    return workflowToken;
   }
 
   async findByToken(token: string) {
@@ -26,6 +104,13 @@ export class WorkflowTokenService {
     return await this.workflowTokenRepository.findByTokenWithExpiredUnscoped(token);
   }
 
+  async count(
+    args: Parameters<typeof this.workflowTokenRepository.count>[0],
+    projectId: TProjectId,
+  ) {
+    return await this.workflowTokenRepository.count(args, projectId);
+  }
+
   async deleteByToken(token: string) {
     return await this.workflowTokenRepository.deleteByTokenUnscoped(token);
   }

services/workflows-service/src/collection-flow/collection-flow.service.intg.test.ts

@@ -37,6 +37,7 @@ import { Customer, EndUser, PrismaClient, Project } from '@prisma/client';
 import { noop } from 'lodash';
 import { CollectionFlowService } from './collection-flow.service';
 import { MerchantMonitoringClient } from '@/business-report/merchant-monitoring-client';
+import { env } from '@/env';
 
 const deps: Provider[] = [
   {
@@ -121,6 +122,7 @@ describe('CollectionFlowService', () => {
   let workflowRuntimeDataRepository: WorkflowRuntimeDataRepository;
   let customerRepository: CustomerRepository;
   let endUserRepository: EndUserRepository;
+  let uiDefinitionRepository: UiDefinitionRepository;
 
   let customer: Customer;
   let project: Project;
@@ -159,6 +161,7 @@ describe('CollectionFlowService', () => {
     );
     customerRepository = module.get<CustomerRepository>(CustomerRepository);
     endUserRepository = module.get<EndUserRepository>(EndUserRepository);
+    uiDefinitionRepository = module.get<UiDefinitionRepository>(UiDefinitionRepository);
   });
 
   beforeEach(async () => {
@@ -202,6 +205,16 @@ describe('CollectionFlowService', () => {
         },
       });
 
+      await uiDefinitionRepository.create({
+        data: {
+          uiSchema: {},
+          projectId: project.id,
+          name: 'test-ui-definition',
+          uiContext: 'collection_flow',
+          workflowDefinitionId: workflowDefinition.id,
+        },
+      });
+
       const workflowRuntimeData = await workflowRuntimeDataRepository.create({
         data: {
           workflowDefinitionId: workflowDefinition.id,
@@ -220,7 +233,28 @@ describe('CollectionFlowService', () => {
 
       const context = await collectionFlowService.getCollectionFlowContext(token);
 
-      expect(context.context).toEqual(workflowContext);
+      const expectedContext = {
+        metadata: {
+          token: token.token,
+          webUiSDKUrl: env.WEB_UI_SDK_URL,
+          collectionFlowUrl: env.COLLECTION_FLOW_URL,
+        },
+        collectionFlow: {
+          state: {
+            steps: [],
+            status: 'pending',
+            currentStep: '',
+          },
+          config: {
+            apiUrl: env.APP_API_URL,
+          },
+          additionalInformation: {
+            customerCompany: customer.displayName,
+          },
+        },
+      };
+
+      expect(context.context).toEqual(expectedContext);
       expect(context.config).toEqual(workflowConfig);
     });
   });