Add WorkspaceAuthContextMiddleware (#17487)

## Context Introduces a middleware that automatically sets the workspace auth context in AsyncLocalStorage for HTTP requests, making it available throughout the request lifecycle without explicit parameter passing. The motivation behind this change is to reduce boilerplate and simplify the developer experience when working with workspace data in HTTP request handlers. The Problem (Before) Every HTTP request handler that needed to access workspace data had to: - Extract auth-related info from decorators (@AuthWorkspace(), @AuthUserWorkspaceId(), etc.) in controller/resolver and pass down to services - Build or pass the authContext explicitly (sometimes with type assertion which was flaky) Then call executeInWorkspaceContext(authContext, async () => { ... }) ## Changes - Add WorkspaceAuthContextMiddleware that extracts auth context from the request and stores it in AsyncLocalStorage - Register middleware for GraphQL, metadata, and REST routes (runs after hydration middlewares) - Simplify executeInWorkspaceContext signature: fn is now the first parameter, authContext is optional second - If authContext is not provided, it's automatically retrieved from the storage (set by middleware) - Update all callers (~120 files) to use the new parameter order - Fixes a bug in search where system auth context was used, bypassing RLS feature.
2026-01-27 17:24:51 +00:00
parent dd98146c99
commit 2daebc6d0f
151 changed files with 3743 additions and 4002 deletions
@@ -39,32 +39,29 @@ export class RunWorkflowJob {
  }: RunWorkflowJobData): Promise<void> {
    const authContext = buildSystemAuthContext(workspaceId);

-    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
-      async () => {
-        try {
-          if (lastExecutedStepId) {
-            await this.resumeWorkflowExecution({
-              workspaceId,
-              workflowRunId,
-              lastExecutedStepId,
-            });
-          } else {
-            await this.startWorkflowExecution({
-              workflowRunId,
-              workspaceId,
-            });
-          }
-        } catch (error) {
-          await this.workflowRunWorkspaceService.endWorkflowRun({
+    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(async () => {
+      try {
+        if (lastExecutedStepId) {
+          await this.resumeWorkflowExecution({
            workspaceId,
            workflowRunId,
-            status: WorkflowRunStatus.FAILED,
-            error: error.message,
+            lastExecutedStepId,
+          });
+        } else {
+          await this.startWorkflowExecution({
+            workflowRunId,
+            workspaceId,
          });
        }
-      },
-    );
+      } catch (error) {
+        await this.workflowRunWorkspaceService.endWorkflowRun({
+          workspaceId,
+          workflowRunId,
+          status: WorkflowRunStatus.FAILED,
+          error: error.message,
+        });
+      }
+    }, authContext);
  }

  private async startWorkflowExecution({
@@ -76,11 +76,9 @@ export class WorkflowCleanWorkflowRunsJob {
    const schemaName = getWorkspaceSchemaName(workspaceId);
    const authContext = buildSystemAuthContext(workspaceId);

-    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
-      async () => {
-        const workflowRunsToDelete = await this.coreDataSource.query(
-          `
+    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(async () => {
+      const workflowRunsToDelete = await this.coreDataSource.query(
+        `
          WITH ranked_runs AS (
            SELECT id,
                   ROW_NUMBER() OVER (
@@ -95,23 +93,22 @@ export class WorkflowCleanWorkflowRunsJob {
          WHERE rn > ${NUMBER_OF_WORKFLOW_RUNS_TO_KEEP}
             OR "createdAt" < NOW() - INTERVAL '14 days';
        `,
+      );
+
+      const workflowRunRepository =
+        await this.globalWorkspaceOrmManager.getRepository(
+          workspaceId,
+          WorkflowRunWorkspaceEntity,
+          { shouldBypassPermissionChecks: true },
        );

-        const workflowRunRepository =
-          await this.globalWorkspaceOrmManager.getRepository(
-            workspaceId,
-            WorkflowRunWorkspaceEntity,
-            { shouldBypassPermissionChecks: true },
-          );
+      for (const workflowRunToDelete of workflowRunsToDelete) {
+        await workflowRunRepository.delete(workflowRunToDelete.id);
+      }

-        for (const workflowRunToDelete of workflowRunsToDelete) {
-          await workflowRunRepository.delete(workflowRunToDelete.id);
-        }
-
-        this.logger.log(
-          `Deleted ${workflowRunsToDelete.length} workflow runs for workspace ${workspaceId}`,
-        );
-      },
-    );
+      this.logger.log(
+        `Deleted ${workflowRunsToDelete.length} workflow runs for workspace ${workspaceId}`,
+      );
+    }, authContext);
  }
 }
@@ -51,41 +51,38 @@ export class WorkflowHandleStaledRunsWorkspaceService {
  private async handleStaledRunsForWorkspace(workspaceId: string) {
    const authContext = buildSystemAuthContext(workspaceId);

-    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
-      async () => {
-        const workflowRunRepository =
-          await this.globalWorkspaceOrmManager.getRepository(
-            workspaceId,
-            WorkflowRunWorkspaceEntity,
-            { shouldBypassPermissionChecks: true },
-          );
-
-        const oneHourAgo = new Date(Date.now() - 60 * 60 * 1000);
-
-        const staledWorkflowRuns = await workflowRunRepository.find({
-          where: {
-            status: WorkflowRunStatus.ENQUEUED,
-            enqueuedAt: Or(LessThan(oneHourAgo), IsNull()),
-          },
-        });
-
-        if (staledWorkflowRuns.length <= 0) {
-          return;
-        }
-
-        await workflowRunRepository.update(
-          staledWorkflowRuns.map((workflowRun) => workflowRun.id),
-          {
-            enqueuedAt: null,
-            status: WorkflowRunStatus.NOT_STARTED,
-          },
-        );
-
-        await this.workflowThrottlingWorkspaceService.recomputeWorkflowRunNotStartedCount(
+    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(async () => {
+      const workflowRunRepository =
+        await this.globalWorkspaceOrmManager.getRepository(
          workspaceId,
+          WorkflowRunWorkspaceEntity,
+          { shouldBypassPermissionChecks: true },
        );
-      },
-    );
+
+      const oneHourAgo = new Date(Date.now() - 60 * 60 * 1000);
+
+      const staledWorkflowRuns = await workflowRunRepository.find({
+        where: {
+          status: WorkflowRunStatus.ENQUEUED,
+          enqueuedAt: Or(LessThan(oneHourAgo), IsNull()),
+        },
+      });
+
+      if (staledWorkflowRuns.length <= 0) {
+        return;
+      }
+
+      await workflowRunRepository.update(
+        staledWorkflowRuns.map((workflowRun) => workflowRun.id),
+        {
+          enqueuedAt: null,
+          status: WorkflowRunStatus.NOT_STARTED,
+        },
+      );
+
+      await this.workflowThrottlingWorkspaceService.recomputeWorkflowRunNotStartedCount(
+        workspaceId,
+      );
+    }, authContext);
  }
 }
@@ -48,7 +48,6 @@ export class WorkflowRunEnqueueWorkspaceService {
      const authContext = buildSystemAuthContext(workspaceId);

      await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-        authContext,
        async () => {
          const workflowRunRepository =
            await this.globalWorkspaceOrmManager.getRepository(
@@ -152,6 +151,7 @@ export class WorkflowRunEnqueueWorkspaceService {
            );
          }
        },
+        authContext,
      );
    } catch (error) {
      this.metricsService.incrementCounter({
@@ -80,7 +80,6 @@ export class WorkflowThrottlingWorkspaceService {

    const currentlyNotStartedWorkflowRunCount =
      await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-        authContext,
        async () => {
          const workflowRunRepository =
            await this.globalWorkspaceOrmManager.getRepository(
@@ -95,6 +94,7 @@ export class WorkflowThrottlingWorkspaceService {
            },
          });
        },
+        authContext,
      );

    await this.setWorkflowRunNotStartedCount(
@@ -113,7 +113,6 @@ export class WorkflowThrottlingWorkspaceService {
    const authContext = buildSystemAuthContext(workspaceId);

    return this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
      async () => {
        const workflowRunRepository =
          await this.globalWorkspaceOrmManager.getRepository(
@@ -128,6 +127,7 @@ export class WorkflowThrottlingWorkspaceService {
          },
        });
      },
+      authContext,
    );
  }

@@ -53,38 +53,35 @@ export class DeleteWorkflowRunsCommand extends ActiveOrSuspendedWorkspacesMigrat
  }: RunOnWorkspaceArgs): Promise<void> {
    const authContext = buildSystemAuthContext(workspaceId);

-    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
-      async () => {
-        try {
-          const workflowRunRepository =
-            await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
-              workspaceId,
-              'workflowRun',
-              { shouldBypassPermissionChecks: true },
-            );
-
-          const createdAtCondition = {
-            createdAt: LessThan(
-              this.createdBeforeDate || new Date().toISOString(),
-            ),
-          };
-
-          const workflowRunCount = await workflowRunRepository.count({
-            where: createdAtCondition,
-          });
-
-          if (!options.dryRun && workflowRunCount > 0) {
-            await workflowRunRepository.delete(createdAtCondition);
-          }
-
-          this.logger.log(
-            `${options.dryRun ? ' (DRY RUN): ' : ''}Deleted ${workflowRunCount} workflow runs`,
+    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(async () => {
+      try {
+        const workflowRunRepository =
+          await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
+            workspaceId,
+            'workflowRun',
+            { shouldBypassPermissionChecks: true },
          );
-        } catch (error) {
-          this.logger.error('Error while deleting workflowRun', error);
+
+        const createdAtCondition = {
+          createdAt: LessThan(
+            this.createdBeforeDate || new Date().toISOString(),
+          ),
+        };
+
+        const workflowRunCount = await workflowRunRepository.count({
+          where: createdAtCondition,
+        });
+
+        if (!options.dryRun && workflowRunCount > 0) {
+          await workflowRunRepository.delete(createdAtCondition);
        }
-      },
-    );
+
+        this.logger.log(
+          `${options.dryRun ? ' (DRY RUN): ' : ''}Deleted ${workflowRunCount} workflow runs`,
+        );
+      } catch (error) {
+        this.logger.error('Error while deleting workflowRun', error);
+      }
+    }, authContext);
  }
 }
@@ -57,7 +57,6 @@ export class WorkflowRunWorkspaceService {
    const authContext = buildSystemAuthContext(workspaceId);

    return this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
      async () => {
        const workflowRunRepository =
          await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
@@ -136,6 +135,7 @@ export class WorkflowRunWorkspaceService {

        return workflowRun.id;
      },
+      authContext,
    );
  }

@@ -346,7 +346,6 @@ export class WorkflowRunWorkspaceService {
    const authContext = buildSystemAuthContext(workspaceId);

    return this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
      async () => {
        const workflowRunRepository =
          await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
@@ -359,6 +358,7 @@ export class WorkflowRunWorkspaceService {
          where: { id: workflowRunId },
        });
      },
+      authContext,
    );
  }

@@ -395,35 +395,32 @@ export class WorkflowRunWorkspaceService {
  }) {
    const authContext = buildSystemAuthContext(workspaceId);

-    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(
-      authContext,
-      async () => {
-        const workflowRunRepository =
-          await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
-            workspaceId,
-            'workflowRun',
-            { shouldBypassPermissionChecks: true },
-          );
-
-        const workflowRunToUpdate = await workflowRunRepository.findOneBy({
-          id: workflowRunId,
-        });
-
-        if (!workflowRunToUpdate) {
-          throw new WorkflowRunException(
-            `workflowRun ${workflowRunId} not found`,
-            WorkflowRunExceptionCode.WORKFLOW_RUN_NOT_FOUND,
-          );
-        }
-
-        await workflowRunRepository.update(
-          workflowRunToUpdate.id,
-          partialUpdate,
-          undefined,
-          ['id'],
+    await this.globalWorkspaceOrmManager.executeInWorkspaceContext(async () => {
+      const workflowRunRepository =
+        await this.globalWorkspaceOrmManager.getRepository<WorkflowRunWorkspaceEntity>(
+          workspaceId,
+          'workflowRun',
+          { shouldBypassPermissionChecks: true },
        );
-      },
-    );
+
+      const workflowRunToUpdate = await workflowRunRepository.findOneBy({
+        id: workflowRunId,
+      });
+
+      if (!workflowRunToUpdate) {
+        throw new WorkflowRunException(
+          `workflowRun ${workflowRunId} not found`,
+          WorkflowRunExceptionCode.WORKFLOW_RUN_NOT_FOUND,
+        );
+      }
+
+      await workflowRunRepository.update(
+        workflowRunToUpdate.id,
+        partialUpdate,
+        undefined,
+        ['id'],
+      );
+    }, authContext);
  }

  private getInitState(