feat(cloud-security): add connectionId query parameter for run status

Author: Marfuen

apps/api/src/cloud-security/cloud-security.controller.ts

@@ -3,6 +3,7 @@ import {
   Post,
   Get,
   Param,
+  Query,
   Headers,
   Logger,
   HttpException,
@@ -11,7 +12,10 @@ import {
 } from '@nestjs/common';
 import { HybridAuthGuard } from '../auth/hybrid-auth.guard';
 import { OrganizationId } from '../auth/auth-context.decorator';
-import { CloudSecurityService } from './cloud-security.service';
+import {
+  CloudSecurityService,
+  ConnectionNotFoundError,
+} from './cloud-security.service';
 
 @Controller({ path: 'cloud-security', version: '1' })
 export class CloudSecurityController {
@@ -85,14 +89,26 @@ export class CloudSecurityController {
   @UseGuards(HybridAuthGuard)
   async getRunStatus(
     @Param('runId') runId: string,
+    @Query('connectionId') connectionId: string,
     @OrganizationId() organizationId: string,
   ) {
+    if (!connectionId) {
+      throw new HttpException(
+        'connectionId query parameter is required',
+        HttpStatus.BAD_REQUEST,
+      );
+    }
+
     try {
       return await this.cloudSecurityService.getRunStatus(
         runId,
+        connectionId,
         organizationId,
       );
     } catch (error) {
+      if (error instanceof ConnectionNotFoundError) {
+        throw new HttpException('Connection not found', HttpStatus.NOT_FOUND);
+      }
       const message =
         error instanceof Error ? error.message : 'Failed to get run status';
       throw new HttpException(message, HttpStatus.INTERNAL_SERVER_ERROR);

apps/api/src/cloud-security/cloud-security.service.ts

@@ -29,19 +29,16 @@ export interface ScanResult {
   error?: string;
 }
 
+export class ConnectionNotFoundError extends Error {
+  constructor() {
+    super('Connection not found');
+  }
+}
+
 @Injectable()
 export class CloudSecurityService {
   private readonly logger = new Logger(CloudSecurityService.name);
 
-  // Track which organization owns each trigger.dev run for authorization
-  private readonly runOwnership = new Map<
-    string,
-    { organizationId: string; createdAt: number }
-  >();
-
-  // Clean up stale entries older than 10 minutes
-  private readonly RUN_OWNERSHIP_TTL_MS = 10 * 60 * 1000;
-
   constructor(
     private readonly credentialVaultService: CredentialVaultService,
     private readonly oauthCredentialsService: OAuthCredentialsService,
@@ -259,49 +256,36 @@ export class CloudSecurityService {
       runId: handle.id,
     });
 
-    // Track ownership for authorization on status checks
-    this.runOwnership.set(handle.id, {
-      organizationId,
-      createdAt: Date.now(),
-    });
-    this.cleanupStaleRuns();
-
     return { runId: handle.id };
   }
 
   async getRunStatus(
     runId: string,
+    connectionId: string,
     organizationId: string,
   ): Promise<{ completed: boolean; success: boolean; output: unknown }> {
-    // Verify the caller's organization owns this run
-    const ownership = this.runOwnership.get(runId);
-    if (!ownership || ownership.organizationId !== organizationId) {
-      throw new Error('Run not found');
+    // Verify the connection belongs to the caller's organization
+    const connection = await db.integrationConnection.findFirst({
+      where: {
+        id: connectionId,
+        organizationId,
+      },
+      select: { id: true },
+    });
+
+    if (!connection) {
+      throw new ConnectionNotFoundError();
     }
 
     const run = await runs.retrieve(runId);
 
-    // Clean up completed runs from the ownership map
-    if (run.isCompleted) {
-      this.runOwnership.delete(runId);
-    }
-
     return {
       completed: run.isCompleted,
       success: run.isCompleted ? run.isSuccess : false,
       output: run.isCompleted ? run.output : null,
     };
   }
 
-  private cleanupStaleRuns(): void {
-    const now = Date.now();
-    for (const [runId, entry] of this.runOwnership) {
-      if (now - entry.createdAt > this.RUN_OWNERSHIP_TTL_MS) {
-        this.runOwnership.delete(runId);
-      }
-    }
-  }
-
   private async storeFindings(
     connectionId: string,
     provider: string,

apps/app/src/app/(app)/[orgId]/cloud-tests/actions/run-platform-scan.ts

@@ -89,7 +89,7 @@ export const runPlatformScan = async (connectionId: string) => {
           provider?: string;
           scannedAt?: string;
         } | null;
-      }>(`/v1/cloud-security/runs/${runId}`, authHeaders);
+      }>(`/v1/cloud-security/runs/${runId}?connectionId=${connectionId}`, authHeaders);
 
       if (statusResponse.error) {
         return {