v6.1.0 - /generate and QOL improvements

.gitignore

@@ -157,3 +157,5 @@ entrypoint.js
 index-K4K5CPUZ.html
 node_modules
 saves/*
+
+.trigger

agent/events.js

@@ -1,5 +1,4 @@
 const { EventEmitter2 } = require("eventemitter2");
-const { censorSensitiveDataDeep } = require("./lib/censorship");
 
 // Factory function to create a new emitter instance with censoring middleware
 const createEmitter = () => {
@@ -13,14 +12,6 @@ const createEmitter = () => {
     ignoreErrors: false,
   });
 
-  // Override emit to censor sensitive data before emitting
-  const originalEmit = emitter.emit.bind(emitter);
-  emitter.emit = function (event, ...args) {
-    // Censor all arguments passed to emit
-    const censoredArgs = args.map(censorSensitiveDataDeep);
-    return originalEmit(event, ...censoredArgs);
-  };
-
   return emitter;
 };
 
@@ -46,7 +37,7 @@ const events = {
   status: "status",
   log: {
     markdown: {
-      static: "log:markdown:static",
+      static: "log:markdown",
       start: "log:markdown:start",
       chunk: "log:markdown:chunk",
       end: "log:markdown:end",

agent/index.js

@@ -63,7 +63,11 @@ class TestDriverAgent extends EventEmitter2 {
     // Derive properties from cliArgs
     const flags = cliArgs.options || {};
     const firstArg = cliArgs.args && cliArgs.args[0];
+
+    // All commands (run, edit, generate) use the same pattern:
+    // first argument is the main file to work with
     this.thisFile = firstArg || this.config.TD_DEFAULT_TEST_FILE;
+
     this.resultFile = flags.resultFile || null;
     this.newSandbox = flags.newSandbox || false;
     this.healMode = flags.healMode || flags.heal || false;
@@ -427,6 +431,7 @@ class TestDriverAgent extends EventEmitter2 {
 
     // Log current execution position for debugging
     if (this.sourceMapper.currentFileSourceMap) {
+      this.emitter.emit(events.log.log, "");
       this.emitter.emit(
         events.log.log,
         theme.dim(`${this.sourceMapper.getCurrentPositionDescription()}`),
@@ -886,30 +891,30 @@ commands:
   // based on the current state of the system (primarily the current screenshot)
   // it will generate files that contain only "prompts"
   // @todo revit the generate command
-  async generate(type, count, baseYaml, skipYaml = false) {
-    this.emitter.emit(events.log.debug, "generate called, %s", type);
+  async generate(count = 1) {
+    this.emitter.emit(events.log.debug, `generate called with count: ${count}`);
 
-    this.emitter.emit(events.log.narration, theme.dim("thinking..."), true);
+    await this.runLifecycle("prerun");
 
-    if (baseYaml && !skipYaml) {
-      await this.runLifecycle("prerun");
-      await this.run(baseYaml, false, false);
-      await this.runLifecycle("postrun");
-    }
+    this.emitter.emit(events.log.narration, theme.dim("thinking..."), true);
 
     let image = await this.system.captureScreenBase64();
 
     const streamId = `generate-${Date.now()}`;
     this.emitter.emit(events.log.markdown.start, streamId);
 
+    let mouse = await this.system.getMousePosition();
+    let activeWindow = await this.system.activeWin();
+
     let message = await this.sdk.req(
       "generate",
       {
-        type,
+        prompt: "make sure to do a spellcheck",
         image,
-        mousePosition: await this.system.getMousePosition(),
-        activeWindow: await this.system.activeWin(),
+        mousePosition: mouse,
+        activeWindow: activeWindow,
         count,
+        stream: false,
       },
       (chunk) => {
         if (chunk.type === "data") {
@@ -932,35 +937,36 @@ commands:
           .replace(/['"`]/g, "")
           .replace(/[^a-zA-Z0-9-]/g, "") // remove any non-alphanumeric chars except hyphens
           .toLowerCase() + ".yaml";
+
       let path1 = path.join(
         this.workingDir,
         "testdriver",
         "generate",
         fileName,
       );
-
       // create generate directory if it doesn't exist
-      if (!fs.existsSync(path.join(this.workingDir, "generate"))) {
-        fs.mkdirSync(path.join(this.workingDir, "generate"));
+      const generateDir = path.join(this.workingDir, "testdriver", "generate");
+      if (!fs.existsSync(generateDir)) {
+        fs.mkdirSync(generateDir);
+        console.log("Created generate directory:", generateDir);
+      } else {
+        console.log("Generate directory already exists:", generateDir);
       }
 
       let list = testPrompt.steps;
 
-      if (baseYaml && fs.existsSync(baseYaml)) {
-        list.unshift({
-          step: {
-            command: "run",
-            file: baseYaml,
-          },
-        });
-      }
       let contents = yaml.dump({
         version: packageJson.version,
         steps: list,
       });
+
+      this.emitter.emit(events.log.debug, `writing file ${path1} ${contents}`);
+
       fs.writeFileSync(path1, contents);
     }
 
+    await this.runLifecycle("postrun");
+
     this.exit(false);
   }
 
@@ -1511,6 +1517,8 @@ ${regression}
   }
 
   async embed(file, depth, pushToHistory) {
+    let inputFile = JSON.parse(JSON.stringify(file));
+
     this.analytics.track("embed", { file });
 
     this.emitter.emit(
@@ -1520,7 +1528,7 @@ ${regression}
 
     depth = depth + 1;
 
-    this.emitter.emit(events.log.log, `${file} (start)`);
+    this.emitter.emit(events.log.log, `${inputFile} (start)`);
 
     // Use the new helper method to resolve file paths relative to testdriver directory
     const currentFilePath = this.sourceMapper.currentFilePath || this.thisFile;
@@ -1573,7 +1581,7 @@ ${regression}
       this.sourceMapper.restoreContext(previousContext);
     }
 
-    this.emitter.emit(events.log.log, `${file} (end)`);
+    this.emitter.emit(events.log.log, `${inputFile} (end)`);
   }
 
   // Returns sandboxId to use (either from file if recent, or null)
@@ -1780,17 +1788,14 @@ ${regression}
       // Start the debugger server as early as possible to ensure event listeners are attached
       if (!debuggerStarted) {
         debuggerStarted = true; // Prevent multiple starts, especially when running test in parallel
-        this.emitter.emit(
-          events.log.narration,
-          theme.green(`Starting debugger server...`),
-        );
         debuggerProcess = await createDebuggerProcess(
           this.config,
           this.emitter,
         );
       }
       this.debuggerUrl = debuggerProcess.url || null; // Store the debugger URL
       this.emitter.emit(events.log.log, `This is beta software!`);
+      this.emitter.emit(events.log.log, ``);
       this.emitter.emit(
         events.log.log,
         theme.yellow(`Join our Discord for help`),
@@ -1799,6 +1804,7 @@ ${regression}
         events.log.log,
         `https://discord.com/invite/cWDFW8DzPm`,
       );
+      this.emitter.emit(events.log.log, ``);
 
       // make testdriver directory if it doesn't exist
       let testdriverFolder = path.join(this.workingDir);
@@ -1812,7 +1818,10 @@ ${regression}
       }
 
       // if the directory for thisFile doesn't exist, create it
-      if (this.cliArgs.command !== "sandbox") {
+      if (
+        this.cliArgs.command !== "sandbox" &&
+        this.cliArgs.command !== "generate"
+      ) {
         const dir = path.dirname(this.thisFile);
         if (!fs.existsSync(dir)) {
           fs.mkdirSync(dir, { recursive: true });
@@ -1837,7 +1846,10 @@ ${regression}
         await this.sdk.auth();
       }
 
-      if (this.cliArgs.command !== "sandbox") {
+      if (
+        this.cliArgs.command !== "sandbox" &&
+        this.cliArgs.command !== "generate"
+      ) {
         this.emitter.emit(
           events.log.log,
           theme.dim(`Working on ${this.thisFile}`),
@@ -2034,6 +2046,20 @@ Please check your network connection, TD_API_KEY, or the service status.`,
     // Use the current file path from sourceMapper to find the lifecycle directory
     // If sourceMapper doesn't have a current file, use thisFile which should be the file being run
     let currentFilePath = this.sourceMapper.currentFilePath || this.thisFile;
+
+    this.emitter.emit(events.log.log, ``);
+    this.emitter.emit(events.log.log, "Running lifecycle: " + lifecycleName);
+
+    // If we still don't have a currentFilePath, fall back to the default testdriver directory
+    if (!currentFilePath) {
+      currentFilePath = path.join(
+        this.workingDir,
+        "testdriver",
+        "testdriver.yaml",
+      );
+      console.log("No currentFilePath found, using fallback:", currentFilePath);
+    }
+
     // Ensure we have an absolute path
     if (currentFilePath && !path.isAbsolute(currentFilePath)) {
       currentFilePath = path.resolve(this.workingDir, currentFilePath);
@@ -2070,6 +2096,9 @@ Please check your network connection, TD_API_KEY, or the service status.`,
         }
       }
     }
+
+    this.emitter.emit(events.log.log, lifecycleFile);
+
     if (lifecycleFile) {
       // Store current source mapping state before running lifecycle file
       const previousContext = this.sourceMapper.saveContext();
@@ -2139,7 +2168,7 @@ Please check your network connection, TD_API_KEY, or the service status.`,
     }
 
     // Move environment setup and special handling here
-    if (["edit", "run"].includes(commandName)) {
+    if (["edit", "run", "generate"].includes(commandName)) {
       await this.buildEnv(options);
     }
 

agent/interface.js

@@ -196,6 +196,42 @@ function createCommandDefinitions(agent) {
         console.log(`TestDriver.ai v${packageJson.version}`);
       },
     },
+
+    generate: {
+      description: "Generate test files based on current screen state",
+      args: {
+        file: Args.string({
+          description: "Base test file to run before generating (optional)",
+          required: false,
+        }),
+      },
+      flags: {
+        count: Flags.integer({
+          description: "Number of test files to generate",
+          default: 3,
+        }),
+        headless: Flags.boolean({
+          description: "Run in headless mode (no GUI)",
+          default: false,
+        }),
+        new: Flags.boolean({
+          description:
+            "Create a new sandbox instead of reconnecting to an existing one",
+          default: false,
+        }),
+        "sandbox-ami": Flags.string({
+          description: "Specify AMI ID for sandbox instance (e.g., ami-1234)",
+        }),
+        "sandbox-instance": Flags.string({
+          description: "Specify EC2 instance type for sandbox (e.g., i3.metal)",
+        }),
+      },
+      handler: async (args, flags) => {
+        // The file argument is already handled by thisFile in the agent constructor
+        // Just call generate with the count
+        await agent.generate(flags.count || 3);
+      },
+    },
   };
 }
 

agent/lib/censorship.js

@@ -38,18 +38,23 @@ const censorSensitiveData = (message) => {
 
 // Function to censor sensitive data in any value (recursive for objects/arrays)
 const censorSensitiveDataDeep = (value) => {
-  if (typeof value === "string") {
-    return censorSensitiveData(value);
-  } else if (Array.isArray(value)) {
-    return value.map(censorSensitiveDataDeep);
-  } else if (value && typeof value === "object") {
-    const result = {};
-    for (const [key, val] of Object.entries(value)) {
-      result[key] = censorSensitiveDataDeep(val);
+  try {
+    if (typeof value === "string") {
+      return censorSensitiveData(value);
+    } else if (Array.isArray(value)) {
+      return value.map(censorSensitiveDataDeep);
+    } else if (value && typeof value === "object") {
+      const result = {};
+      for (const [key, val] of Object.entries(value)) {
+        result[key] = censorSensitiveDataDeep(val);
+      }
+      return result;
     }
-    return result;
+    return value;
+  } catch {
+    // If we hit any error (like circular reference), just return a safe placeholder
+    return "[Object]";
   }
-  return value;
 };
 
 // Function to update interpolation variables (for runtime updates)