From c2d34316b867a5f3aaad405187512138d7359e3f Mon Sep 17 00:00:00 2001
From: Zubair <zubair@smythos.com>
Date: Tue, 1 Jul 2025 00:11:24 +0500
Subject: [PATCH] fixed serverless code component vault key fetching issue

---
 .../src/Components/ServerlessCode.class.ts    | 20 ++++++++-
 .../core/src/helpers/AWSLambdaCode.helper.ts  | 43 ++++++++++++++++++-
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/Components/ServerlessCode.class.ts b/packages/core/src/Components/ServerlessCode.class.ts
index 7bb8973d..d2772475 100644
--- a/packages/core/src/Components/ServerlessCode.class.ts
+++ b/packages/core/src/Components/ServerlessCode.class.ts
@@ -2,6 +2,8 @@ import { IAgent as Agent } from '@sre/types/Agent.types';
 import { Component } from './Component.class';
 import Joi from 'joi';
 import { ConnectorService } from '@sre/Core/ConnectorsService';
+import { AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
+import { calculateExecutionCost, getLambdaCredentials, reportUsage } from '@sre/helpers/AWSLambdaCode.helper';
 
 export class ServerlessCode extends Component {
 
@@ -56,9 +58,18 @@ export class ServerlessCode extends Component {
             }
 
             logger.debug(`\nInput Variables: \n${JSON.stringify(codeInputs, null, 2)}\n`);
-            
+
             let codeConnector = ConnectorService.getCodeConnector();
+            let codeCredentials: AWSCredentials & AWSRegionConfig & { isUserProvidedKeys: boolean } =
+                await getLambdaCredentials(agent, config);
 
+            if (codeCredentials.isUserProvidedKeys) {
+                codeConnector = codeConnector.instance({
+                    region: codeCredentials.region,
+                    accessKeyId: codeCredentials.accessKeyId,
+                    secretAccessKey: codeCredentials.secretAccessKey,
+                })
+            }
             // Deploy lambda function if it doesn't exist or the code hash is different
             await codeConnector.agent(agent.id)
                 .deploy(config.id, {
@@ -78,6 +89,13 @@ export class ServerlessCode extends Component {
                 );
                 logger.debug(`Execution time: ${executionTime}ms\n`);
 
+                const cost = calculateExecutionCost(executionTime);
+                if (!codeCredentials.isUserProvidedKeys) {
+                    const accountConnector = ConnectorService.getAccountConnector();
+                    const agentTeam = await accountConnector.getCandidateTeam(agent.id);
+                    reportUsage({ cost, agentId: agent.id, teamId: agentTeam });
+                }
+
                 if (executionResponse.success) {
                     Output = executionResponse.output;
                 } else {
diff --git a/packages/core/src/helpers/AWSLambdaCode.helper.ts b/packages/core/src/helpers/AWSLambdaCode.helper.ts
index 85bf7dcd..178bdfb6 100644
--- a/packages/core/src/helpers/AWSLambdaCode.helper.ts
+++ b/packages/core/src/helpers/AWSLambdaCode.helper.ts
@@ -5,9 +5,13 @@ import zl from 'zip-lib';
 import { InvokeCommand, Runtime, LambdaClient, UpdateFunctionCodeCommand, CreateFunctionCommand, GetFunctionCommand, GetFunctionCommandOutput, InvokeCommandOutput } from '@aws-sdk/client-lambda';
 import { GetRoleCommand, CreateRoleCommand, IAMClient, GetRoleCommandOutput, CreateRoleCommandOutput } from '@aws-sdk/client-iam';
 import fs from 'fs';
-import { AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
+import { AWSConfig, AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
+import { VaultHelper } from '@sre/Security/Vault.service/Vault.helper';
+import { IAgent } from '@sre/types/Agent.types';
+import { SystemEvents } from '@sre/Core/SystemEvents';
 export const cachePrefix = 'serverless_code';
 export const cacheTTL = 60 * 60 * 24 * 16; // 16 days
+const PER_SECOND_COST = 0.0001;
 
 export function getLambdaFunctionName(agentId: string, componentId: string) {
     return `${agentId}-${componentId}`;
@@ -340,8 +344,45 @@ export async function getDeployedFunction(functionName: string, awsConfigs: AWSC
     }
 }
 
+export async function getLambdaCredentials(agent: IAgent, config: any): Promise<AWSConfig & { isUserProvidedKeys: boolean }> {
+    let awsAccessKeyId = null;
+    let awsSecretAccessKey = null;
+    let awsRegion = null;
+    let userProvidedKeys = false;
+    if (config.data.accessKeyId && config.data.secretAccessKey && config.data.region && config.data.use_own_keys) {
+        userProvidedKeys = true;
+        [awsAccessKeyId, awsSecretAccessKey] = await Promise.all([
+            VaultHelper.getTeamKey(extractKeyFromTemplateVar(config.data.accessKeyId), agent?.teamId),
+            VaultHelper.getTeamKey(extractKeyFromTemplateVar(config.data.secretAccessKey), agent?.teamId),
+        ]);
+        awsRegion = config.data.region;
+    }
+    const awsCredentials = {
+        accessKeyId: awsAccessKeyId,
+        secretAccessKey: awsSecretAccessKey,
+        region: awsRegion,
+        isUserProvidedKeys: userProvidedKeys,
+    };
+    return awsCredentials;
+}
+
+export function calculateExecutionCost(executionTime: number) {
+    // executionTime in milliseconds
+    const cost = (executionTime / 1000) * Number(PER_SECOND_COST);
+    return cost;
+}
+
 export function extractKeyFromTemplateVar(input: string) {
     const regex = /\{\{KEY\((.*?)\)\}\}/;
     const match = input.match(regex);
     return match ? match[1] : input;
+}
+
+export function reportUsage({ cost, agentId, teamId }: { cost: number; agentId: string; teamId: string }) {
+    SystemEvents.emit('USAGE:API', {
+        sourceId: 'api:serverless_code.smyth',
+        cost,
+        agentId,
+        teamId,
+    });
 }
\ No newline at end of file