agentlang 0.0.5 → 0.0.7

package/src/runtime/auth/cognito.ts

@@ -6,14 +6,7 @@ import {
   SignUpCallback,
   UserInfo,
 } from './interface.js';
-import {
-  ensureUser,
-  ensureUserSession,
-  findUser,
-  findUserByEmail,
-  findUserSession,
-  removeSession,
-} from '../modules/auth.js';
+import { ensureUser, ensureUserSession, findUser, findUserByEmail } from '../modules/auth.js';
 import { logger } from '../logger.js';
 import { sleepMilliseconds } from '../util.js';
 import { Instance } from '../module.js';
@@ -36,9 +29,14 @@ let fromEnv: any = undefined;
 let CognitoIdentityProviderClient: any = undefined;
 let SignUpCommand: any = undefined;
 let AdminGetUserCommand: any = undefined;
+let InitiateAuthCommand: any = undefined;
 let AuthenticationDetails: any = undefined;
 let CognitoUser: any = undefined;
 let CognitoUserPool: any = undefined;
+let CognitoUserSession: any = undefined;
+let CognitoIdToken: any = undefined;
+let CognitoAccessToken: any = undefined;
+let CognitoRefreshToken: any = undefined;
 
 if (isNodeEnv) {
   const cp = await import('@aws-sdk/credential-providers');
@@ -48,11 +46,16 @@ if (isNodeEnv) {
   CognitoIdentityProviderClient = cip.CognitoIdentityProviderClient;
   SignUpCommand = cip.SignUpCommand;
   AdminGetUserCommand = cip.AdminGetUserCommand;
+  InitiateAuthCommand = cip.InitiateAuthCommand;
 
   const ci = await import('amazon-cognito-identity-js');
   AuthenticationDetails = ci.AuthenticationDetails;
   CognitoUser = ci.CognitoUser;
   CognitoUserPool = ci.CognitoUserPool;
+  CognitoUserSession = ci.CognitoUserSession;
+  CognitoIdToken = ci.CognitoIdToken;
+  CognitoAccessToken = ci.CognitoAccessToken;
+  CognitoRefreshToken = ci.CognitoRefreshToken;
 }
 
 const defaultConfig = isNodeEnv
@@ -441,7 +444,13 @@ export class CognitoAuth implements AgentlangAuth {
         const idToken = result.getIdToken().getJwtToken();
         const accessToken = result.getAccessToken().getJwtToken();
         const refreshToken = result.getRefreshToken().getToken();
-        const localSess: Instance = await ensureUserSession(userid, idToken, env);
+        const localSess: Instance = await ensureUserSession(
+          userid,
+          idToken,
+          accessToken,
+          refreshToken,
+          env
+        );
         const sessInfo: SessionInfo = {
           sessionId: localSess.lookup('id'),
           userId: userid,
@@ -513,7 +522,13 @@ export class CognitoAuth implements AgentlangAuth {
         const idToken = result.getIdToken().getJwtToken();
         const accessToken = result.getAccessToken().getJwtToken();
         const refreshToken = result.getRefreshToken().getToken();
-        const localSess: Instance = await ensureUserSession(userid, idToken, env);
+        const localSess: Instance = await ensureUserSession(
+          userid,
+          idToken,
+          accessToken,
+          refreshToken,
+          env
+        );
         const sessInfo: SessionInfo = {
           sessionId: localSess.lookup('id'),
           userId: userid,
@@ -540,25 +555,33 @@ export class CognitoAuth implements AgentlangAuth {
         return;
       }
       const user = new CognitoUser({
-        Username: localUser.email,
+        Username: localUser.lookup('email'),
         Pool: this.fetchUserPool(),
       });
+
       let done = false;
       let logoutError: any;
-      user.signOut(() => {
-        done = true;
+
+      const session = new CognitoUserSession({
+        IdToken: new CognitoIdToken({ IdToken: sessionInfo.idToken }),
+        AccessToken: new CognitoAccessToken({ AccessToken: sessionInfo.accessToken }),
+        RefreshToken: new CognitoRefreshToken({ RefreshToken: sessionInfo.refreshToken }),
       });
-      // Add error handling for signOut
-      user.signOut((err: any) => {
-        if (err) {
+      user.setSignInUserSession(session);
+      user.globalSignOut({
+        onSuccess: function () {
+          done = true;
+        },
+        onFailure: function (err: any) {
+          done = true;
           logger.error(`Cognito signOut error for user ${sessionInfo.userId}:`, {
             errorName: err.name,
             errorMessage: sanitizeErrorMessage(err.message),
           });
           logoutError = err;
-        }
-        done = true;
+        },
       });
+
       while (!done) {
         await sleepMilliseconds(100);
       }
@@ -568,10 +591,6 @@ export class CognitoAuth implements AgentlangAuth {
         );
         // Continue with local session cleanup even if Cognito logout fails
       }
-      const sess = await findUserSession(localUser.id, env);
-      if (sess) {
-        await removeSession(sess.id, env);
-      }
       logger.debug(`Successfully logged out user ${sessionInfo.userId}`);
       if (cb) cb(true);
     } catch (err: any) {
@@ -581,6 +600,50 @@ export class CognitoAuth implements AgentlangAuth {
     }
   }
 
+  async changePassword(
+    sessionInfo: SessionInfo,
+    newPassword: string,
+    oldPassword: string,
+    env: Environment
+  ): Promise<boolean> {
+    const localUser = await findUser(sessionInfo.userId, env);
+    if (!localUser) {
+      logger.warn(`User ${sessionInfo.userId} not found for password-change`);
+      return false;
+    }
+    const email = localUser.lookup('email');
+    const user = new CognitoUser({
+      Username: email,
+      Pool: this.fetchUserPool(),
+    });
+    const session = new CognitoUserSession({
+      IdToken: new CognitoIdToken({ IdToken: sessionInfo.idToken }),
+      AccessToken: new CognitoAccessToken({ AccessToken: sessionInfo.accessToken }),
+      RefreshToken: new CognitoRefreshToken({ RefreshToken: sessionInfo.refreshToken }),
+    });
+    user.setSignInUserSession(session);
+    let done = false;
+    let cpErr: any = undefined;
+    user.changePassword(oldPassword, newPassword, (err: any, _: any) => {
+      if (err) {
+        done = true;
+        cpErr = err;
+      } else {
+        done = true;
+      }
+    });
+
+    while (!done) {
+      await sleepMilliseconds(100);
+    }
+
+    if (cpErr) {
+      logger.warn(`Failed to change the password for ${email} - ${cpErr.message}`);
+      return false;
+    }
+    return true;
+  }
+
   private fetchUserPool() {
     if (this.userPool) {
       return this.userPool;
@@ -753,4 +816,71 @@ export class CognitoAuth implements AgentlangAuth {
       };
     }
   }
+
+  async refreshToken(refreshTokenString: string, env: Environment): Promise<SessionInfo> {
+    try {
+      // Use InitiateAuth with REFRESH_TOKEN_AUTH flow
+      const client = new CognitoIdentityProviderClient({
+        region: process.env.AWS_REGION || 'us-west-2',
+        credentials: fromEnv(),
+      });
+
+      const command = new InitiateAuthCommand({
+        AuthFlow: 'REFRESH_TOKEN_AUTH',
+        ClientId: this.fetchClientId(),
+        AuthParameters: {
+          REFRESH_TOKEN: refreshTokenString,
+        },
+      });
+
+      const response = await client.send(command);
+
+      if (!response.AuthenticationResult) {
+        throw new UnauthorisedError('Token refresh failed');
+      }
+
+      const newIdToken = response.AuthenticationResult.IdToken!;
+      const newAccessToken = response.AuthenticationResult.AccessToken!;
+      const newRefreshToken = response.AuthenticationResult.RefreshToken || refreshTokenString;
+
+      // Extract user info from the new ID token
+      const idTokenPayload = JSON.parse(atob(newIdToken.split('.')[1]));
+      const userEmail = idTokenPayload.email;
+
+      // Find or create local user
+      let localUser = await findUserByEmail(userEmail, env);
+      if (!localUser) {
+        localUser = await ensureUser(userEmail, '', '', env);
+      }
+      const userId = localUser.lookup('id');
+
+      // Update local session
+      const updatedSession = await ensureUserSession(
+        userId,
+        newIdToken,
+        newAccessToken,
+        newRefreshToken,
+        env
+      );
+
+      const sessInfo: SessionInfo = {
+        sessionId: updatedSession.lookup('id'),
+        userId: userId,
+        authToken: newIdToken,
+        idToken: newIdToken,
+        accessToken: newAccessToken,
+        refreshToken: newRefreshToken,
+        systemSesionInfo: response.AuthenticationResult,
+      };
+
+      return sessInfo;
+    } catch (err: any) {
+      logger.error(`Refresh token operation failed: ${err.message}`);
+      if (err.name === 'NotAuthorizedException') {
+        throw new UnauthorisedError('Invalid or expired refresh token');
+      }
+      handleCognitoError(err, 'refreshToken');
+      throw err; // This line won't be reached due to handleCognitoError throwing
+    }
+  }
 }

package/src/runtime/auth/interface.ts

@@ -33,4 +33,11 @@ export interface AgentlangAuth {
   verifyToken(token: string, env?: Environment): any;
   getUser(userId: string, env: Environment): Promise<UserInfo>;
   getUserByEmail(email: string, env: Environment): Promise<UserInfo>;
+  changePassword(
+    sessionInfo: SessionInfo,
+    newPassword: string,
+    oldPassword: string,
+    env: Environment
+  ): Promise<boolean>;
+  refreshToken(refreshToken: string, env: Environment): Promise<SessionInfo>;
 }

package/src/runtime/interpreter.ts

@@ -54,7 +54,6 @@ import {
   escapeFqName,
   escapeQueryName,
   fqNameFromPath,
-  invokeModuleFn,
   isFqName,
   isPath,
   isString,
@@ -78,6 +77,7 @@ import {
   maybeCancelTimer,
   setTimerRunning,
 } from './modules/core.js';
+import { invokeModuleFn } from './jsmodules.js';
 
 export type Result = any;
 
@@ -590,8 +590,8 @@ export async function parseAndEvaluateStatement(
 }
 
 async function evaluatePattern(pat: Pattern, env: Environment): Promise<void> {
-  if (pat.literal) {
-    await evaluateLiteral(pat.literal, env);
+  if (pat.expr) {
+    await evaluateExpression(pat.expr, env);
   } else if (pat.crudMap) {
     await evaluateCrudMap(pat.crudMap, env);
   } else if (pat.forEach) {
@@ -1047,21 +1047,61 @@ async function walkJoinQueryPattern(
 async function handleAgentInvocation(agentEventInst: Instance, env: Environment): Promise<void> {
   const agent: AgentInstance = await findAgentByName(agentEventInst.name, env);
   await agent.invoke(agentEventInst.lookup('message'), env);
-  const result: string = env.getLastResult();
-  if (agent.isPlanner()) {
-    logger.debug(`Agent ${agent.name} generated pattern: ${result}`);
-    try {
-      if (result.trimStart().startsWith('workflow')) {
-        await parseWorkflow(result);
-        return;
-      } else {
-        env.setLastResult(await parseAndEvaluateStatement(result, undefined, env));
+  const result: string | undefined = cleanupAgentResponse(env.getLastResult());
+  if (result) {
+    if (agent.isPlanner()) {
+      logger.debug(`Agent ${agent.name} generated pattern: ${result}`);
+      try {
+        let rs = result.trim();
+        if (rs.startsWith('[')) {
+          const stmts = rs.substring(1, rs.length - 1);
+          rs = `workflow T {${stmts}}`;
+        }
+        if (rs.startsWith('workflow')) {
+          const wf = await parseWorkflow(rs);
+          if (agent.runWorkflows) {
+            await evaluateStatements(wf.statements, env);
+          }
+        } else {
+          env.setLastResult(await parseAndEvaluateStatement(rs, undefined, env));
+        }
+      } catch (err: any) {
+        logger.error(
+          `Failed to evaluate pattern generated by agent ${agent.name} - ${result}, ${err}`
+        );
       }
-    } catch (err: any) {
-      logger.error(
-        `Failed to evaluate pattern generated by agent ${agent.name} - ${result}, ${err}`
-      );
     }
+  } else {
+    logger.warn(`Agent ${agent.name} failed to generate a response`);
+  }
+}
+
+function cleanupAgentResponse(response: string | undefined): string | undefined {
+  if (response) {
+    const resp = response.trim();
+    if (resp.startsWith('[') && resp.endsWith(']')) {
+      return resp;
+    }
+    const parts = resp.split('\n');
+    const validated = parts.filter((s: string) => {
+      const stmt = s.trim();
+      const r =
+        stmt.startsWith('{') ||
+        stmt.startsWith('}') ||
+        stmt.startsWith('if') ||
+        stmt.startsWith('for') ||
+        stmt.startsWith('delete') ||
+        stmt.startsWith('workflow');
+      if (!r) {
+        const i = stmt.indexOf('(');
+        return i > 0 && stmt.indexOf(')') > i;
+      } else {
+        return r;
+      }
+    });
+    return validated.join('\n');
+  } else {
+    return response;
   }
 }
 
@@ -1263,7 +1303,12 @@ async function applyFn(fnCall: FnCall, env: Environment, isAsync: boolean): Prom
     if (fnCall.args != undefined) {
       args = new Array<Result>();
       for (let i = 0; i < fnCall.args.length; ++i) {
-        await evaluateLiteral(fnCall.args[i], env);
+        const arg = fnCall.args[i];
+        if (isLiteral(arg)) {
+          await evaluateLiteral(arg, env);
+        } else {
+          await evaluateExpression(arg, env);
+        }
         args.push(env.getLastResult());
       }
       args.push(env);
@@ -1287,7 +1332,7 @@ async function realizeMap(mapLiteral: MapLiteral, env: Environment): Promise<voi
   for (let i = 0; i < mapLiteral.entries.length; ++i) {
     const entry = mapLiteral.entries[i];
     const k = getMapKey(entry.key);
-    await evaluateLiteral(entry.value, env);
+    await evaluateExpression(entry.value, env);
     result.set(k, env.getLastResult());
   }
   env.setLastResult(Object.fromEntries(result.entries()));

package/src/runtime/jsmodules.ts

@@ -0,0 +1,123 @@
+import { logger } from './logger.js';
+import { setSubscription } from './resolvers/registry.js';
+import { now, splitRefs } from './util.js';
+
+const importedModules = new Map<string, any>();
+
+// Usage: importModule("./mymodels/acme.js")
+export async function importModule(path: string, name: string) {
+  if (!(path.startsWith('/') || path.startsWith('.'))) {
+    path = process.cwd() + '/' + path;
+  }
+  const m = await import(/* @vite-ignore */ path);
+  importedModules.set(name, m);
+  // e.g of dynamic fn-call:
+  //// let f = eval("(a, b) => m.add(a, b)");
+  //// console.log(f(10, 20))
+  return m;
+}
+
+export function moduleImported(moduleName: string): boolean {
+  return importedModules.has(moduleName);
+}
+
+const ReservedImports = new Set<string>(['resolvers']);
+
+export function valiadteImportName(n: string) {
+  if (ReservedImports.has(n)) {
+    throw new Error(`${n} is an import reserved by the runtime`);
+  }
+}
+
+function maybeEvalFunction(fnName: string): Function | undefined {
+  try {
+    return eval(fnName);
+  } catch (reason: any) {
+    logger.debug(reason);
+    return undefined;
+  }
+}
+
+function invokeReservedFn(moduleName: string, fnName: string, args: Array<any> | null) {
+  if (moduleName == 'resolvers' && fnName == 'setSubscription' && args != null) {
+    return setSubscription(args[0], args[1]);
+  } else {
+    throw new Error(`Failed to call ${moduleName}.${fnName} with the given arguments`);
+  }
+}
+
+async function invokeBuiltInFn(fnName: string, args: Array<any> | null, isAsync: boolean = false) {
+  if (fnName == 'now') {
+    return now();
+  } else if (fnName == 'uuid') {
+    return crypto.randomUUID();
+  } else {
+    const pf: Function | undefined = maybeEvalFunction(fnName);
+    if (pf instanceof Function) {
+      if (args == null) {
+        if (isAsync) return await pf();
+        else return pf();
+      } else {
+        if (isAsync) return await pf(...args.slice(0, args.length - 1));
+        else return pf(...args.slice(0, args.length - 1));
+      }
+    } else {
+      throw new Error(`Failed to invoke function - ${fnName}`);
+    }
+  }
+}
+
+export async function invokeModuleFn(
+  fqFnName: string,
+  args: Array<any> | null,
+  isAsync: boolean = false
+): Promise<any> {
+  try {
+    const refs: string[] = splitRefs(fqFnName);
+    if (refs.length == 1) {
+      if (isAsync) {
+        return await invokeBuiltInFn(refs[0], args, isAsync);
+      } else {
+        return invokeBuiltInFn(refs[0], args);
+      }
+    }
+    const mname = refs[0];
+    if (ReservedImports.has(mname)) {
+      return invokeReservedFn(mname, refs[1], args);
+    }
+    const m = importedModules.get(mname);
+    if (m != undefined) {
+      const f = m[refs[1]];
+      if (f != undefined) {
+        if (args == null)
+          if (isAsync) {
+            return await f();
+          } else return f();
+        else if (isAsync) {
+          return await f(...args);
+        } else return f(...args);
+      } else throw new Error(`Function not found - ${fqFnName}`);
+    } else throw new Error(`JavaScript module ${refs[0]} not found`);
+  } catch (reason: any) {
+    const pf: Function | undefined = maybeEvalFunction(fqFnName);
+    if (pf instanceof Function) {
+      if (args == null) {
+        if (isAsync) return await pf();
+        else return pf();
+      } else {
+        if (isAsync) return await pf(...args.slice(0, args.length - 1));
+        else return pf(...args.slice(0, args.length - 1));
+      }
+    } else {
+      throw new Error(reason);
+    }
+  }
+}
+
+export function getModuleFn(fqFnName: string): Function | undefined {
+  const refs: string[] = splitRefs(fqFnName);
+  const m = importedModules.get(refs[0]);
+  if (m != undefined) {
+    return m[refs[1]];
+  } else return undefined;
+}

package/src/runtime/loader.ts

@@ -23,6 +23,7 @@ import {
   ResolverDefinition,
   ResolverMethodSpec,
   AgentPropertyDef,
+  isLiteral,
 } from '../language/generated/ast.js';
 import {
   addEntity,
@@ -45,14 +46,11 @@ import {
 } from './module.js';
 import {
   findRbacSchema,
-  getModuleFn,
-  importModule,
   isString,
   makeFqName,
   maybeExtends,
   registerInitFunction,
   runShellCommand,
-  //runShellCommand,
 } from './util.js';
 import { getFileSystem, toFsPath, readFile, readdir, exists } from '../utils/fs-utils.js';
 import { URI } from 'vscode-uri';
@@ -67,6 +65,7 @@ import { AgentEntityName, CoreAIModuleName, LlmEntityName } from './modules/ai.j
 import { GenericResolver, GenericResolverMethods } from './resolvers/interface.js';
 import { registerResolver, setResolver, setSubscription } from './resolvers/registry.js';
 import { ConfigSchema } from './state.js';
+import { getModuleFn, importModule, valiadteImportName } from './jsmodules.js';
 
 export async function extractDocument(
   fileName: string,
@@ -403,7 +402,7 @@ export async function runStandaloneStatements() {
           GlobalEnvironment.switchActiveModuleName(oldModule);
         }
       }
-      logger.info(`Init eval result: ${GlobalEnvironment.getLastResult().toString()}`);
+      logger.debug(`Init eval result: ${GlobalEnvironment.getLastResult()}`);
     });
     StandaloneStatements.clear();
   }
@@ -421,10 +420,12 @@ async function addAgentDefinition(def: AgentDefinition, moduleName: string) {
     if (apdef.value.array) {
       v = apdef.value.array.vals
         .map((stmt: Statement) => {
-          if (stmt.pattern.literal) {
-            const s = stmt.pattern.literal.str;
+          if (stmt.pattern.expr && isLiteral(stmt.pattern.expr)) {
+            const s = stmt.pattern.expr.str || stmt.pattern.expr.id;
             if (s == undefined) {
-              throw new Error(`Only arrays of string-literals are to be passed to agent ${name}`);
+              throw new Error(
+                `Only arrays of string-literals or identifiers should be passed to agent ${name}`
+              );
             }
             return s;
           } else {
@@ -499,12 +500,11 @@ function addResolverDefinition(def: ResolverDefinition, moduleName: string) {
     paths.forEach((path: string) => {
       setResolver(path, resolverName);
     });
-    if (subsFn && subsEvent) {
+    if (subsFn) {
       resolver.subs = {
         subscribe: subsFn,
-        onSubscriptionEvent: subsEvent,
       };
-      setSubscription(subsEvent, resolverName);
+      if (subsEvent) setSubscription(subsEvent, resolverName);
       resolver.subscribe();
     }
   });
@@ -549,6 +549,7 @@ export async function internModule(module: ModuleDefinition): Promise<Module> {
   const mn = module.name;
   const r = addModule(mn);
   module.imports.forEach(async (imp: Import) => {
+    valiadteImportName(imp.name);
     await importModule(imp.path, imp.name);
   });
   for (let i = 0; i < module.defs.length; ++i) {