Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions packages/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@
"build:types": "tsc --emitDeclarationOnly --declaration --outDir dist/types -p tsconfig.dts.json",
"build:jsbundle": "cross-env rollup -c",
"build": "pnpm run build:jsbundle && pnpm run build:types",
"test:unit": "cd ../.. && vitest run packages/core/tests/unit",
"test:unit:watch": "cd ../.. && vitest watch packages/core/tests/unit",
"test:integration": "cd ../.. && vitest run packages/core/tests/integration",
"test:integration:watch": "cd ../.. && vitest watch packages/core/tests/integration",
"gen:docs": "typedoc",
"doc:graphgen": "npx depcruise src --config ./doc/.dep-minimal.json --output-type dot > ./doc/dep-graph.dot && dot -Tpng ./doc/dep-graph.dot -o ./doc/dep-graph.png",
"knip": "knip"
Expand Down
7 changes: 4 additions & 3 deletions packages/core/src/Components/ECMASandbox.class.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,15 @@ export class ECMASandbox extends Component {
}
}

const inputVarsCode = this.generateInputVarCode(codeInputs);
const code = inputVarsCode + '\n' + config.data.code;
//const inputVarsCode = this.generateInputVarCode(codeInputs);
//const code = inputVarsCode + '\n' + config.data.code;
const code = config.data.code;

logger.debug(`Running code: \n${code}\n`);

const ecmaCodeConnector = ConnectorService.getCodeConnector('ECMASandbox');

const executionResponse: CodeExecutionResult = await ecmaCodeConnector.agent(agent.id).execute(config.id, { code });
const executionResponse: CodeExecutionResult = await ecmaCodeConnector.agent(agent.id).execute(config.id, { code, inputs: input });
if (executionResponse.success) {
Output = executionResponse.output;
} else {
Expand Down
52 changes: 30 additions & 22 deletions packages/core/src/helpers/AWSLambdaCode.helper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,17 @@ import crypto from 'crypto';
import { ConnectorService } from '@sre/Core/ConnectorsService';
import { AccessCandidate } from '@sre/Security/AccessControl/AccessCandidate.class';
import zl from 'zip-lib';
import { InvokeCommand, Runtime, LambdaClient, UpdateFunctionCodeCommand, CreateFunctionCommand, GetFunctionCommand, GetFunctionCommandOutput, InvokeCommandOutput, UpdateFunctionConfigurationCommand } from '@aws-sdk/client-lambda';
import {
InvokeCommand,
Runtime,
LambdaClient,
UpdateFunctionCodeCommand,
CreateFunctionCommand,
GetFunctionCommand,
GetFunctionCommandOutput,
InvokeCommandOutput,
UpdateFunctionConfigurationCommand,
} from '@aws-sdk/client-lambda';
import { GetRoleCommand, CreateRoleCommand, IAMClient, GetRoleCommandOutput, CreateRoleCommandOutput } from '@aws-sdk/client-iam';
import fs from 'fs';
import { AWSConfig, AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
Expand All @@ -19,7 +29,6 @@ export function getLambdaFunctionName(agentId: string, componentId: string) {
return `${agentId}-${componentId}`;
}


export function generateCodeHash(code_body: string, codeInputs: string[], envVariables: string[]) {
const bodyHash = getSanitizeCodeHash(code_body);
const inputsHash = getSanitizeCodeHash(JSON.stringify(codeInputs));
Expand Down Expand Up @@ -79,9 +88,7 @@ export async function getDeployedCodeHash(agentId: string, componentId: string)

export async function setDeployedCodeHash(agentId: string, componentId: string, codeHash: string) {
const redisCache = ConnectorService.getCacheConnector();
await redisCache
.user(AccessCandidate.agent(agentId))
.set(`${cachePrefix}_${agentId}-${componentId}`, codeHash, null, null, cacheTTL);
await redisCache.user(AccessCandidate.agent(agentId)).set(`${cachePrefix}_${agentId}-${componentId}`, codeHash, null, null, cacheTTL);
}

function replaceVaultKeysTemplateVars(code: string, envVariables: Record<string, string>) {
Expand Down Expand Up @@ -121,7 +128,7 @@ export async function zipCode(directory: string) {
},
function (err) {
reject(err);
},
}
);
});
}
Expand Down Expand Up @@ -276,7 +283,6 @@ export function getLambdaRolePolicy() {
});
}


export async function updateDeployedCodeTTL(agentId: string, componentId: string, ttl: number) {
const redisCache = ConnectorService.getCacheConnector();
await redisCache.user(AccessCandidate.agent(agentId)).updateTTL(`${cachePrefix}_${agentId}-${componentId}`, ttl);
Expand All @@ -285,7 +291,7 @@ export async function updateDeployedCodeTTL(agentId: string, componentId: string
export async function invokeLambdaFunction(
functionName: string,
inputs: { [key: string]: any },
awsCredentials: AWSCredentials & AWSRegionConfig,
awsCredentials: AWSCredentials & AWSRegionConfig
): Promise<any> {
try {
const client = new LambdaClient({
Expand Down Expand Up @@ -384,11 +390,11 @@ export function reportUsage({ cost, agentId, teamId }: { cost: number; agentId:

export function validateAsyncMainFunction(rawCode: string): { isValid: boolean; error?: string; parameters?: string[]; dependencies?: string[] } {
try {
const code = replaceVaultKeysTemplateVars(rawCode, {});
const code = replaceVaultKeysTemplateVars(rawCode.trim(), {});
// Parse the code using acorn
const ast = acorn.parse(code, {
ecmaVersion: 'latest',
sourceType: 'module'
sourceType: 'module',
});

// Extract library imports
Expand All @@ -414,22 +420,26 @@ export function validateAsyncMainFunction(rawCode: string): { isValid: boolean;
}

// Handle CallExpression (require() calls)
if (node.type === 'CallExpression' &&
if (
node.type === 'CallExpression' &&
node.callee.type === 'Identifier' &&
node.callee.name === 'require' &&
node.arguments.length > 0 &&
node.arguments[0].type === 'Literal') {
node.arguments[0].type === 'Literal'
) {
const modulePath = node.arguments[0].value;
if (modulePath && !modulePath.startsWith('.') && !modulePath.startsWith('/')) {
libraries.add(extractPackageName(modulePath));
}
}

// Handle dynamic import() calls
if (node.type === 'CallExpression' &&
if (
node.type === 'CallExpression' &&
node.callee.type === 'Import' &&
node.arguments.length > 0 &&
node.arguments[0].type === 'Literal') {
node.arguments[0].type === 'Literal'
) {
const modulePath = node.arguments[0].value;
if (modulePath && !modulePath.startsWith('.') && !modulePath.startsWith('/')) {
libraries.add(extractPackageName(modulePath));
Expand Down Expand Up @@ -503,23 +513,23 @@ export function validateAsyncMainFunction(rawCode: string): { isValid: boolean;
return {
isValid: false,
error: 'No main function found at root level',
dependencies
dependencies,
};
}

if (!hasAsyncMain) {
return {
isValid: false,
error: 'Main function exists but is not async',
dependencies
dependencies,
};
}

return { isValid: true, parameters: mainParameters, dependencies };
} catch (error) {
return {
isValid: false,
error: `Failed to parse code: ${error.message}`
error: `Failed to parse code: ${error.message}`,
};
}
}
Expand Down Expand Up @@ -549,7 +559,7 @@ export function generateCodeFromLegacyComponent(code_body: string, code_imports:
async function main(${codeInputs.join(', ')}) {
${code_body}
}
`
`;
return code;
}

Expand All @@ -565,14 +575,12 @@ export function extractAllKeyNamesFromTemplateVars(input: string): string[] {
return matches;
}


async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string, key: string }> {
async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string; key: string }> {
const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
return {
value: vaultSecret,
key: keyName,
};

}

export async function getCurrentEnvironmentVariables(agentTeamId: string, code: string): Promise<Record<string, string>> {
Expand All @@ -588,4 +596,4 @@ export async function getCurrentEnvironmentVariables(agentTeamId: string, code:
export function getSortedObjectValues(obj: Record<string, string>): string[] {
const sortedKeys = Object.keys(obj).sort();
return sortedKeys.map((key) => obj[key]);
}
}
5 changes: 3 additions & 2 deletions packages/core/src/helpers/Conversation.helper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -132,9 +132,9 @@ export class Conversation extends EventEmitter {
console.warn('Conversation Error: ', error?.message);
});
this._maxContextSize =
_settings.maxContextSize || (this._model as TLLMModel).tokens || (this._model as TLLMModel).keyOptions?.tokens || this._maxContextSize;
_settings?.maxContextSize || (this._model as TLLMModel).tokens || (this._model as TLLMModel).keyOptions?.tokens || this._maxContextSize;
this._maxOutputTokens =
_settings.maxOutputTokens ||
_settings?.maxOutputTokens ||
(this._model as TLLMModel).completionTokens ||
(this._model as TLLMModel).keyOptions?.completionTokens ||
this._maxOutputTokens;
Expand Down Expand Up @@ -928,6 +928,7 @@ export class Conversation extends EventEmitter {
//is this a valid agent data?
if (typeof specSource?.behavior === 'string' && specSource?.components && specSource?.connections) {
this.agentData = specSource; //agent loaded from data directly
this._agentId = specSource.id;
return await this.loadSpecFromAgent(specSource);
}

Expand Down
24 changes: 17 additions & 7 deletions packages/core/src/helpers/ECMASandbox.helper.ts
Original file line number Diff line number Diff line change
@@ -1,11 +1,22 @@
import 'ses';




export function runJs(code: string) {
// Call lockdown to secure the environment
lockdown();
// Ensure SES lockdown happens only once per process
// and tolerate environments already locked down elsewhere
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const globalAny: any = globalThis as any;
if (globalAny.__SRE_SES_LOCKED_DOWN__ !== true) {
try {
lockdown();
} catch (err) {
const msg = String(err || '');
if (!(msg.includes('Already locked down') || msg.includes('SES_ALREADY_LOCKED_DOWN'))) {
throw err;
}
} finally {
globalAny.__SRE_SES_LOCKED_DOWN__ = true;
}
}
try {
// Endow the compartment with necessary APIs
const compartment = new Compartment({
Expand All @@ -27,7 +38,6 @@ export function runJs(code: string) {
console.error(error);
throw error;
}

}

function getParametersString(parameters: string[], inputs: Record<string, any>) {
Expand All @@ -49,6 +59,6 @@ export function generateExecutableCode(code: string, parameters: string[], input
const result = await main(${getParametersString(parameters, inputs)});
return result;
})();
`
`;
return executableCode;
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@ import fs from 'fs';
import path from 'path';
import { AgentDataConnector } from '../AgentDataConnector';
import { uid } from '@sre/utils/general.utils';
import { Logger } from '@sre/helpers/Log.helper';

const console = Logger('LocalAgentDataConnector');

export type LocalAgentDataSettings = { devDir: string; prodDir: string };

Expand Down
Loading