feat: support @ embed references in WCS prompts

This is useful for supporting the standard that is currently implemented
by Gemini CLI, and makes it easy to re-use system prompts, or context files
that are originally built for e.g. Gemini CLI

Author: devversion

examples/environments/remote_env/fake-executor.ts

@@ -85,6 +85,10 @@ export class FakeRemoteExecutor implements Executor {
     };
   }
 
+  async executeProjectTests() {
+    return null;
+  }
+
   async shouldRepairFailedBuilds() {
     // Some environments have a builtin retry loop as part of initial generation.
     // In those cases, you may want to skip retrying.

examples/environments/remote_env/other-file-2.md

@@ -0,0 +1 @@
+Also use ABC.

examples/environments/remote_env/other-file.md

@@ -0,0 +1,3 @@
+Use XYZ.
+
+@./other-file-2.md

examples/environments/remote_env/system-instructions.md

@@ -2,3 +2,5 @@ Follow instructions below CAREFULLY:
 
 - Code MUST be implemented in my super secret framework.
 - Put the component code inside `src/app/app.ts`
+
+@./other-file.md

runner/configuration/environment.ts

@@ -1,6 +1,7 @@
 import {readdirSync, readFileSync, statSync} from 'fs';
-import {basename, dirname, extname, join, resolve} from 'path';
+import {basename, extname, join, resolve} from 'path';
 import {globSync} from 'tinyglobby';
+import {Executor} from '../orchestration/executors/executor.js';
 import {Rating} from '../ratings/rating-types.js';
 import {
   FrameworkInfo,
@@ -13,10 +14,7 @@ import {generateId} from '../utils/id-generation.js';
 import {lazy} from '../utils/lazy-creation.js';
 import {EnvironmentConfig} from './environment-config.js';
 import {MultiStepPrompt} from './multi-step-prompt.js';
-import {renderHandlebarsTemplate} from './prompt-templating.js';
-import {RunnerName} from '../codegen/runner-creation.js';
-import {Executor} from '../orchestration/executors/executor.js';
-import {LocalExecutor} from '../orchestration/executors/local-executor.js';
+import {renderPromptTemplate} from './prompt-templating.js';
 
 /** Represents a single prompt evaluation environment. */
 export class Environment {
@@ -133,8 +131,8 @@ export class Environment {
     promptFilePath: string | null,
     additionalContext: Record<string, string> = {},
   ) {
-    return renderHandlebarsTemplate(content, {
-      rootDir: promptFilePath ? dirname(promptFilePath) : null,
+    return renderPromptTemplate(content, {
+      containingFile: promptFilePath ? promptFilePath : null,
       FULL_STACK_FRAMEWORK_NAME: this.fullStackFramework.displayName,
       CLIENT_SIDE_FRAMEWORK_NAME: this.clientSideFramework.displayName,
       ...additionalContext,

runner/configuration/prompt-templating.ts

@@ -1,36 +1,46 @@
-import Handlebars from 'handlebars';
 import {readFileSync} from 'fs';
-import path from 'path';
+import Handlebars from 'handlebars';
+import path, {dirname, resolve} from 'path';
 import {UserFacingError} from '../utils/errors.js';
 
 function initializeHandlebars() {
   Handlebars.registerHelper('neq', (a, b) => a !== b);
-  Handlebars.registerPartial('embed', (ctx: {rootDir: string | null; file?: string}) => {
+  Handlebars.registerPartial('embed', (ctx: {containingFile: string | null; file?: string}) => {
     if (!ctx.file) {
       throw new UserFacingError('file= is required');
     }
-    if (!ctx.rootDir) {
-      throw new UserFacingError('Cannot use `embed` if a rootDir is not specified');
+    if (!ctx.containingFile) {
+      throw new UserFacingError('Cannot use `embed` if not `containingFile` is specified');
     }
 
-    const fullPath = path.join(ctx.rootDir, ctx.file);
-    const content = readFileSync(fullPath, 'utf8');
+    const fullPath = path.join(dirname(ctx.containingFile), ctx.file);
+    let content = readFileSync(fullPath, 'utf8');
+    content = processAtFileReferencesSync(content, ctx.containingFile);
 
     // Recursively support `embed`.
     return Handlebars.compile(content, {strict: true})({
       ...ctx,
-      rootDir: path.dirname(fullPath),
+      containingFile: fullPath,
     });
   });
 }
 
 initializeHandlebars();
 
-/** Renders the given content via Handlebars. */
-export function renderHandlebarsTemplate<T extends {rootDir: string | null}>(
+/**
+ * Renders the given prompt template, by supporting:
+ *   - Handlebars builtins.
+ *   - Handlebars `embed` custom partials for including other files.
+ *   - Supporting the ecosystem `@<file-path>` standard for including other files.
+ *
+ * If `context` does not specify a `containingFile`, then other files cannot be embedded.
+ */
+export function renderPromptTemplate<T extends {containingFile: string | null}>(
   content: string,
   ctx: T,
 ) {
+  content = ctx.containingFile ? processAtFileReferencesSync(content, ctx.containingFile) : content;
+
   const template = Handlebars.compile(content, {strict: true});
   const contextFiles: string[] = [];
   const result = template(ctx, {
@@ -72,3 +82,35 @@ export function renderHandlebarsTemplate<T extends {rootDir: string | null}>(
     contextFiles,
   };
 }
+
+function processAtFileReferencesSync(content: string, containingFile: string): string {
+  let newContent = content;
+  let match;
+  // Match all `@./<file-path>` or `@/<file-path>` occurrences.
+  // If someone intends to write such text in their prompt, they could overcome this
+  // by indenting the string, or adding arbitrary characters before front.
+  const regex = /^@(\.?\/[^\s]+)/gm;
+  const containingFileDir = dirname(containingFile);
+
+  while ((match = regex.exec(newContent)) !== null) {
+    const filePath = match[1];
+    const fullPath = resolve(containingFileDir, filePath);
+    let replacement: string;
+    try {
+      replacement = readFileSync(fullPath, 'utf8');
+      // Note: If we start searching the match start, where the new embedded content begins,
+      // we can trivially. process nested embeds via the `@` syntax.
+    } catch (e) {
+      throw new Error(
+        `Unexpected error while embedding \`${match[0]}\` reference in ${containingFile}. ` +
+          `Error: ${e}`,
+      );
+    }
+
+    newContent =
+      newContent.substring(0, match.index) +
+      processAtFileReferencesSync(replacement, fullPath) +
+      newContent.substring(regex.lastIndex);
+  }
+  return newContent;
+}

runner/eval-cli.ts

@@ -232,7 +232,7 @@ async function handler(cliArgs: Arguments<Options>): Promise<void> {
     } else {
       console.error(chalk.red('An error occurred during the assessment process:'));
       console.error(chalk.red(error));
-      if ((error as Partial<Error>).stack) {
+      if (process.env.DEBUG === '1' && (error as Partial<Error>).stack) {
         console.error(chalk.red((error as Error).stack));
       }
     }