Skip to content

Commit 5c76029

Browse files
Merge pull request #45 from SmythOS/fix/core/lambda-code-connector-refactoring
unified code_imports and code_body into single code key for lambda code connector
2 parents a13d0a7 + 40699cd commit 5c76029

File tree

4 files changed

+208
-58
lines changed

4 files changed

+208
-58
lines changed

packages/core/src/Components/ServerlessCode.class.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,14 @@ import { Component } from './Component.class';
33
import Joi from 'joi';
44
import { ConnectorService } from '@sre/Core/ConnectorsService';
55
import { AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
6-
import { calculateExecutionCost, getLambdaCredentials, reportUsage } from '@sre/helpers/AWSLambdaCode.helper';
6+
import { calculateExecutionCost, generateCodeFromLegacyComponent, getLambdaCredentials, reportUsage } from '@sre/helpers/AWSLambdaCode.helper';
77

88
export class ServerlessCode extends Component {
99

1010
protected configSchema = Joi.object({
1111
code_imports: Joi.string().max(1000).allow('').label('Imports'),
1212
code_body: Joi.string().max(500000).allow('').label('Code'),
13+
code: Joi.string().max(500000).allow('').label('Code').optional(),
1314
deploy_btn: Joi.string().max(500000).allow('').label('Deploy').optional(),
1415
accessKeyId: Joi.string().max(100).allow('').label('AWS Access Key ID').optional(),
1516
secretAccessKey: Joi.string().max(200).allow('').label('AWS Secret Access Key').optional(),
@@ -70,11 +71,14 @@ export class ServerlessCode extends Component {
7071
secretAccessKey: codeCredentials.secretAccessKey,
7172
})
7273
}
74+
let code = config?.data?.code;
75+
if (!code) {
76+
code = generateCodeFromLegacyComponent(config.data.code_body, config.data.code_imports, Object.keys(codeInputs))
77+
}
7378
// Deploy lambda function if it doesn't exist or the code hash is different
7479
await codeConnector.agent(agent.id)
7580
.deploy(config.id, {
76-
code: config?.data?.code_body,
77-
dependencies: config?.data?.code_imports,
81+
code,
7882
inputs: codeInputs,
7983
}, {
8084
runtime: 'nodejs',

packages/core/src/helpers/AWSLambdaCode.helper.ts

Lines changed: 184 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import { AWSConfig, AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types
99
import { VaultHelper } from '@sre/Security/Vault.service/Vault.helper';
1010
import { IAgent } from '@sre/types/Agent.types';
1111
import { SystemEvents } from '@sre/Core/SystemEvents';
12+
import * as acorn from 'acorn';
13+
1214
export const cachePrefix = 'serverless_code';
1315
export const cacheTTL = 60 * 60 * 24 * 16; // 16 days
1416
const PER_SECOND_COST = 0.0001;
@@ -18,11 +20,10 @@ export function getLambdaFunctionName(agentId: string, componentId: string) {
1820
}
1921

2022

21-
export function generateCodeHash(code_body: string, code_imports: string, codeInputs: string[]) {
22-
const importsHash = getSanitizeCodeHash(code_imports);
23+
export function generateCodeHash(code_body: string, codeInputs: string[]) {
2324
const bodyHash = getSanitizeCodeHash(code_body);
2425
const inputsHash = getSanitizeCodeHash(JSON.stringify(codeInputs));
25-
return `imports-${importsHash}__body-${bodyHash}__inputs-${inputsHash}`;
26+
return `body-${bodyHash}__inputs-${inputsHash}`;
2627
}
2728

2829
export function getSanitizeCodeHash(code: string) {
@@ -81,49 +82,17 @@ export async function setDeployedCodeHash(agentId: string, componentId: string,
8182
.set(`${cachePrefix}_${agentId}-${componentId}`, codeHash, null, null, cacheTTL);
8283
}
8384

84-
85-
export function extractNpmImports(code: string) {
86-
const importRegex = /import\s+(?:[\w*\s{},]*\s+from\s+)?['"]([^'"]+)['"]/g;
87-
const requireRegex = /require\(['"]([^'"]+)['"]\)/g;
88-
const dynamicImportRegex = /import\(['"]([^'"]+)['"]\)/g;
89-
90-
let libraries = new Set();
91-
let match;
92-
93-
// Function to extract the main package name
94-
function extractPackageName(modulePath: string) {
95-
if (modulePath.startsWith('@')) {
96-
// Handle scoped packages (e.g., @babel/core)
97-
return modulePath.split('/').slice(0, 2).join('/');
98-
}
99-
return modulePath.split('/')[0]; // Extract the first part (main package)
100-
}
101-
// Match static ESM imports
102-
while ((match = importRegex.exec(code)) !== null) {
103-
libraries.add(extractPackageName(match[1]));
104-
}
105-
// Match CommonJS require() calls
106-
while ((match = requireRegex.exec(code)) !== null) {
107-
libraries.add(extractPackageName(match[1]));
108-
}
109-
// Match dynamic import() calls
110-
while ((match = dynamicImportRegex.exec(code)) !== null) {
111-
libraries.add(extractPackageName(match[1]));
112-
}
113-
114-
return Array.from(libraries);
115-
}
116-
117-
118-
export function generateLambdaCode(code_imports: string, code_body: string, input_variables: string[]) {
119-
const lambdaCode = `${code_imports}\nexport const handler = async (event, context) => {
85+
export function generateLambdaCode(code: string, parameters: string[]) {
86+
const lambdaCode = `
87+
${code}
88+
export const handler = async (event, context) => {
12089
try {
12190
context.callbackWaitsForEmptyEventLoop = false;
12291
let startTime = Date.now();
123-
const result = await (async () => {
124-
${input_variables && input_variables.length ? input_variables.map((variable) => `const ${variable} = event.${variable};`).join('\n') : ''}
125-
${code_body}
126-
})();
92+
93+
${parameters && parameters.length ? parameters.map((variable) => `const ${variable} = event.${variable};`).join('\n') : ''}
94+
const result = await main(${parameters.join(', ')});
95+
12796
let endTime = Date.now();
12897
return {
12998
result,
@@ -385,4 +354,175 @@ export function reportUsage({ cost, agentId, teamId }: { cost: number; agentId:
385354
agentId,
386355
teamId,
387356
});
388-
}
357+
}
358+
359+
export function validateAsyncMainFunction(code: string): { isValid: boolean; error?: string; parameters?: string[]; dependencies?: string[] } {
360+
try {
361+
// Parse the code using acorn
362+
const ast = acorn.parse(code, {
363+
ecmaVersion: 'latest',
364+
sourceType: 'module'
365+
});
366+
367+
// Extract library imports
368+
const libraries = new Set<string>();
369+
function extractPackageName(modulePath: string): string {
370+
if (modulePath.startsWith('@')) {
371+
// Handle scoped packages (e.g., @babel/core)
372+
return modulePath.split('/').slice(0, 2).join('/');
373+
}
374+
return modulePath.split('/')[0]; // Extract the first part (main package)
375+
}
376+
377+
function processNodeForImports(node: any): void {
378+
if (!node) return;
379+
380+
// Handle ImportDeclaration (ES6 imports)
381+
if (node.type === 'ImportDeclaration') {
382+
const modulePath = node.source.value;
383+
if (modulePath && !modulePath.startsWith('.') && !modulePath.startsWith('/')) {
384+
// Skip relative imports and absolute paths
385+
libraries.add(extractPackageName(modulePath));
386+
}
387+
}
388+
389+
// Handle CallExpression (require() calls)
390+
if (node.type === 'CallExpression' &&
391+
node.callee.type === 'Identifier' &&
392+
node.callee.name === 'require' &&
393+
node.arguments.length > 0 &&
394+
node.arguments[0].type === 'Literal') {
395+
const modulePath = node.arguments[0].value;
396+
if (modulePath && !modulePath.startsWith('.') && !modulePath.startsWith('/')) {
397+
libraries.add(extractPackageName(modulePath));
398+
}
399+
}
400+
401+
// Handle dynamic import() calls
402+
if (node.type === 'CallExpression' &&
403+
node.callee.type === 'Import' &&
404+
node.arguments.length > 0 &&
405+
node.arguments[0].type === 'Literal') {
406+
const modulePath = node.arguments[0].value;
407+
if (modulePath && !modulePath.startsWith('.') && !modulePath.startsWith('/')) {
408+
libraries.add(extractPackageName(modulePath));
409+
}
410+
}
411+
412+
// Recursively process child nodes
413+
for (const key in node) {
414+
if (node[key] && typeof node[key] === 'object') {
415+
if (Array.isArray(node[key])) {
416+
(node[key] as any[]).forEach(processNodeForImports);
417+
} else {
418+
processNodeForImports(node[key]);
419+
}
420+
}
421+
}
422+
}
423+
424+
// Extract dependencies from the entire AST
425+
processNodeForImports(ast);
426+
const dependencies = Array.from(libraries) as string[];
427+
428+
// Check if there's a function declaration or function expression named 'main' at the root level
429+
let hasAsyncMain = false;
430+
let hasMain = false;
431+
let mainParameters: string[] = [];
432+
433+
for (const node of ast.body) {
434+
if (node.type === 'FunctionDeclaration') {
435+
if (node.id?.name === 'main') {
436+
hasMain = true;
437+
if (node.async) {
438+
hasAsyncMain = true;
439+
mainParameters = extractParameters(node.params);
440+
break;
441+
}
442+
}
443+
} else if (node.type === 'VariableDeclaration') {
444+
// Check for const/let/var main = async function() or const/let/var main = async () =>
445+
for (const declarator of node.declarations) {
446+
if (declarator.id.type === 'Identifier' && declarator.id.name === 'main') {
447+
hasMain = true;
448+
if (declarator.init) {
449+
if (declarator.init.type === 'FunctionExpression' && declarator.init.async) {
450+
hasAsyncMain = true;
451+
mainParameters = extractParameters(declarator.init.params);
452+
break;
453+
} else if (declarator.init.type === 'ArrowFunctionExpression' && declarator.init.async) {
454+
hasAsyncMain = true;
455+
mainParameters = extractParameters(declarator.init.params);
456+
break;
457+
}
458+
}
459+
}
460+
}
461+
} else if (node.type === 'ExpressionStatement' && node.expression.type === 'AssignmentExpression') {
462+
// Check for main = async function() or main = async () =>
463+
if (node.expression.left.type === 'Identifier' && node.expression.left.name === 'main') {
464+
hasMain = true;
465+
const right = node.expression.right;
466+
if ((right.type === 'FunctionExpression' || right.type === 'ArrowFunctionExpression') && right.async) {
467+
hasAsyncMain = true;
468+
mainParameters = extractParameters(right.params);
469+
break;
470+
}
471+
}
472+
}
473+
}
474+
475+
if (!hasMain) {
476+
return {
477+
isValid: false,
478+
error: 'No main function found at root level',
479+
dependencies
480+
};
481+
}
482+
483+
if (!hasAsyncMain) {
484+
return {
485+
isValid: false,
486+
error: 'Main function exists but is not async',
487+
dependencies
488+
};
489+
}
490+
491+
return { isValid: true, parameters: mainParameters, dependencies };
492+
} catch (error) {
493+
return {
494+
isValid: false,
495+
error: `Failed to parse code: ${error.message}`
496+
};
497+
}
498+
}
499+
500+
function extractParameters(params: any[]): string[] {
501+
return params.map((param: any): string => {
502+
if (param.type === 'Identifier') {
503+
return param.name;
504+
} else if (param.type === 'AssignmentPattern' && param.left.type === 'Identifier') {
505+
return param.left.name;
506+
} else if (param.type === 'RestElement' && param.argument.type === 'Identifier') {
507+
return param.argument.name;
508+
} else if (param.type === 'ObjectPattern') {
509+
// For destructured objects, return the object name or a placeholder
510+
return param.name || '[object]';
511+
} else if (param.type === 'ArrayPattern') {
512+
// For destructured arrays, return a placeholder
513+
return '[array]';
514+
}
515+
return '[unknown]';
516+
});
517+
}
518+
519+
export function generateCodeFromLegacyComponent(code_body: string, code_imports: string, codeInputs: string[]) {
520+
const code = `
521+
${code_imports}
522+
async function main(${codeInputs.join(', ')}) {
523+
${code_body}
524+
}
525+
`
526+
return code;
527+
}
528+

packages/core/src/subsystems/ComputeManager/Code.service/connectors/AWSLambdaCode.class.ts

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { AccessRequest } from '@sre/Security/AccessControl/AccessRequest.class';
55
import fs from 'fs';
66
import path from 'path';
77
import { execSync } from 'child_process';
8-
import { cacheTTL, createOrUpdateLambdaFunction, extractNpmImports, generateCodeHash, generateLambdaCode, getDeployedCodeHash, getDeployedFunction, getLambdaFunctionName, invokeLambdaFunction, setDeployedCodeHash, updateDeployedCodeTTL, zipCode } from '@sre/helpers/AWSLambdaCode.helper';
8+
import { cacheTTL, createOrUpdateLambdaFunction, generateCodeHash, generateLambdaCode, getDeployedCodeHash, getDeployedFunction, getLambdaFunctionName, invokeLambdaFunction, setDeployedCodeHash, updateDeployedCodeTTL, validateAsyncMainFunction, zipCode } from '@sre/helpers/AWSLambdaCode.helper';
99
import { AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
1010
import { Logger } from '@sre/helpers/Log.helper';
1111

@@ -38,7 +38,7 @@ export class AWSLambdaCode extends CodeConnector {
3838
getDeployedCodeHash(agentId, codeUID),
3939
]);
4040

41-
const codeHash = generateCodeHash(input.code, input.dependencies, Object.keys(input.inputs));
41+
const codeHash = generateCodeHash(input.code, Object.keys(input.inputs));
4242
if (isLambdaExists && exisitingCodeHash === codeHash) {
4343
return {
4444
id: functionName,
@@ -53,22 +53,24 @@ export class AWSLambdaCode extends CodeConnector {
5353
}
5454
const directory = `${baseFolder}/${functionName}__${Date.now()}`;
5555
try {
56-
const libraries = extractNpmImports(input.dependencies);
57-
58-
const lambdaCode = generateLambdaCode(input.dependencies, input.code, Object.keys(input.inputs));
56+
const { isValid, parameters, error, dependencies } = validateAsyncMainFunction(input.code);
57+
if (!isValid) {
58+
throw new Error(error || 'Invalid Code')
59+
}
60+
const lambdaCode = generateLambdaCode(input.code, parameters);
5961
// create folder
6062
fs.mkdirSync(directory);
6163
// create index.js file
6264
fs.writeFileSync(path.join(directory, 'index.mjs'), lambdaCode);
6365
// run command npm init
6466
execSync('npm init -y', { cwd: directory });
6567
// run command npm install
66-
execSync(`npm install ${libraries.join(' ')}`, { cwd: directory });
68+
execSync(`npm install ${dependencies.join(' ')}`, { cwd: directory });
6769

6870
const zipFilePath = await zipCode(directory);
6971
await createOrUpdateLambdaFunction(functionName, zipFilePath, this.awsConfigs);
7072
await setDeployedCodeHash(agentId, codeUID, codeHash);
71-
console.log('Lambda function updated successfully!');
73+
console.debug('Lambda function updated successfully!');
7274
return {
7375
id: functionName,
7476
runtime: config.runtime,
@@ -78,8 +80,10 @@ export class AWSLambdaCode extends CodeConnector {
7880
} catch (error) {
7981
throw error;
8082
} finally {
81-
fs.rmSync(`${directory}`, { recursive: true, force: true });
82-
fs.unlinkSync(`${directory}.zip`);
83+
try {
84+
fs.rmSync(`${directory}`, { recursive: true, force: true });
85+
fs.unlinkSync(`${directory}.zip`);
86+
} catch (error) {}
8387
}
8488

8589
}
@@ -93,8 +97,8 @@ export class AWSLambdaCode extends CodeConnector {
9397
const executionTime = lambdaResponse.executionTime;
9498
await updateDeployedCodeTTL(agentId, codeUID, cacheTTL);
9599
console.debug(
96-
`Code result:\n ${typeof lambdaResponse.result === 'object' ? JSON.stringify(lambdaResponse.result, null, 2) : lambdaResponse.result
97-
}\n`,
100+
`Code result:\n ${typeof lambdaResponse.result === 'object' ? JSON.stringify(lambdaResponse.result, null, 2) : lambdaResponse.result
101+
}\n`,
98102
);
99103
console.debug(`Execution time: ${executionTime}ms\n`);
100104

packages/sdk/src/Components/generated/ServerlessCode.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ export interface TServerlessCodeSettings {
1111
code_imports?: string;
1212
/** Code */
1313
code_body?: string;
14+
/** Code */
15+
code?: string;
1416
/** Deploy */
1517
deploy_btn?: string;
1618
/** AWS Access Key ID */

0 commit comments

Comments
 (0)