trace included files

packages/poml/file.tsx

@@ -56,6 +56,7 @@ interface Range {
 export class PomlFile {
   private text: string;
   private sourcePath: string | undefined;
+  private tracePath: string | undefined;
   private config: PomlReaderConfig;
   private ast: XMLDocument | undefined;
   private cst: CstNode;
@@ -67,6 +68,7 @@ export class PomlFile {
   private responseSchema: Schema | undefined;
   private toolsSchema: ToolsSchema | undefined;
   private runtimeParameters: { [key: string]: any } | undefined;
+  private includedTraces: { markup: string; context: any; sourcePath?: string }[] = [];
 
   constructor(text: string, options?: PomlReaderOptions, sourcePath?: string) {
     this.config = {
@@ -76,6 +78,7 @@ export class PomlFile {
     };
     this.text = this.config.crlfToLf ? text.replace(/\r\n/g, '\n') : text;
     this.sourcePath = sourcePath;
+    this.tracePath = sourcePath;
     if (this.sourcePath) {
       const envFile = this.sourcePath.replace(/(source)?\.poml$/i, '.env');
       if (existsSync(envFile)) {
@@ -243,6 +246,10 @@ export class PomlFile {
     return this.runtimeParameters;
   }
 
+  public getIncludedTraces() {
+    return this.includedTraces;
+  }
+
   public xmlRootElement(): XMLElement | undefined {
     if (!this.ast || !this.ast.rootElement) {
       this.reportError('Root element is invalid.', {
@@ -629,22 +636,58 @@ export class PomlFile {
 
     const source = src.value;
 
+    let includePath =
+      this.sourcePath && !path.isAbsolute(source) ? path.join(path.dirname(this.sourcePath), source) : source;
+
+    if (this.tracePath) {
+      const dir = path.dirname(this.tracePath);
+      const base = path.basename(this.tracePath).replace(/\.source\.poml$/i, '.poml');
+      const match = base.match(/^(\d{4})\./);
+      if (match) {
+        const idx = match[1];
+        const envPath = path.join(dir, base.replace(/\.poml$/i, '.env'));
+        if (existsSync(envPath)) {
+          try {
+            const envText = readFileSync(envPath, 'utf8');
+            const mainMatch = envText.match(/^SOURCE_PATH=(.*)$/m);
+            if (mainMatch) {
+              const originalMain = mainMatch[1];
+              const originalInclude = path.isAbsolute(source) ? source : path.join(path.dirname(originalMain), source);
+              const childBase = path.basename(originalInclude, '.poml');
+              const childEnv = path.join(dir, `${idx}.${childBase}.env`);
+              if (existsSync(childEnv)) {
+                const childEnvText = readFileSync(childEnv, 'utf8');
+                const childMatch = childEnvText.match(/^SOURCE_PATH=(.*)$/m);
+                if (childMatch && childMatch[1] === originalInclude) {
+                  const useSource = this.tracePath.endsWith('.source.poml');
+                  const candidate = path.join(dir, `${idx}.${childBase}${useSource ? '.source.poml' : '.poml'}`);
+                  if (existsSync(candidate)) {
+                    includePath = candidate;
+                  }
+                }
+              }
+            }
+          } catch {
+            /* ignore */
+          }
+        }
+      }
+    }
+
     let text: string;
     try {
-      text = readSource(source, this.sourcePath ? path.dirname(this.sourcePath) : undefined, 'string');
+      text = readFileSync(includePath, 'utf8');
     } catch (e) {
       this.reportError(
-        e !== undefined && (e as Error).message ? (e as Error).message : `Error reading source: ${source}`,
+        e !== undefined && (e as Error).message ? (e as Error).message : `Error reading source: ${includePath}`,
         this.xmlAttributeValueRange(src),
         e,
       );
       return <></>;
     }
 
-    const includePath =
-      this.sourcePath && !path.isAbsolute(source) ? path.join(path.dirname(this.sourcePath), source) : source;
-
     const included = new PomlFile(text, this.config, includePath);
+    this.includedTraces.push({ markup: text, context, sourcePath: includePath }, ...included.getIncludedTraces());
     const root = included.xmlRootElement();
     if (!root) {
       return <></>;

packages/poml/index.ts

@@ -18,7 +18,7 @@ import './presentation';
 import './essentials';
 import './components';
 import { reactRender } from './util/reactRender';
-import { dumpTrace, setTrace, clearTrace, isTracing, parseJsonWithBuffers } from './util/trace';
+import { dumpTrace, dumpTraceInclude, setTrace, clearTrace, isTracing, parseJsonWithBuffers } from './util/trace';
 
 export type { RichContent, Message, SourceMapRichContent, SourceMapMessage };
 export { richContentFromSourceMap };
@@ -206,11 +206,7 @@ export async function commandLine(args: CliArgs) {
 
   ErrorCollection.clear();
 
-  const pomlFile = new PomlFile(input, readOptions, sourcePath);
-  let reactElement = pomlFile.react(context);
-  reactElement = React.createElement(StyleSheetProvider, { stylesheet }, reactElement);
-
-  const ir = await read(input, readOptions, context, stylesheet, sourcePath);
+  const [ir, pomlFile] = await _readWithFile(input, readOptions, context, stylesheet, sourcePath);
 
   const speakerMode = args.speakerMode === true || args.speakerMode === undefined;
   const prettyPrint = args.prettyPrint === true;
@@ -222,15 +218,20 @@ export async function commandLine(args: CliArgs) {
     : renderContent(resultMessages as RichContent);
   const result: CliResult = {
     messages: resultMessages,
-    schema: pomlFile.getResponseSchema()?.toOpenAPI(),
-    tools: pomlFile.getToolsSchema()?.toOpenAI(),
-    runtime: pomlFile.getRuntimeParameters(),
+    schema: pomlFile?.getResponseSchema()?.toOpenAPI(),
+    tools: pomlFile?.getToolsSchema()?.toOpenAI(),
+    runtime: pomlFile?.getRuntimeParameters(),
   };
   const output = prettyPrint ? prettyOutput : JSON.stringify(result);
 
   if (isTracing()) {
     try {
-      dumpTrace(input, context, stylesheet, result, sourcePath, prettyOutput);
+      const basePrefix = dumpTrace(input, context, stylesheet, result, sourcePath, prettyOutput);
+      if (basePrefix && pomlFile) {
+        for (const t of pomlFile.getIncludedTraces()) {
+          dumpTraceInclude(basePrefix, t.markup, t.context, t.sourcePath);
+        }
+      }
     } catch (err: any) {
       ErrorCollection.add(new SystemError('Failed to dump trace', { cause: err }));
     }

packages/poml/tests/trace.test.tsx

@@ -66,6 +66,34 @@ describe('trace dumps', () => {
     fs.rmSync(origDir, { recursive: true, force: true });
   });
 
+  test('included files are traced and used when original missing', async () => {
+    const origDir = fs.mkdtempSync(path.join(os.tmpdir(), 'orig-'));
+    const mainPath = path.join(origDir, 'main.poml');
+    const childPath = path.join(origDir, 'includeChild.poml');
+    fs.copyFileSync(path.join(__dirname, 'assets', 'includeChild.poml'), childPath);
+    fs.writeFileSync(mainPath, '<poml><include src="includeChild.poml"/></poml>');
+
+    await commandLine({ file: mainPath, speakerMode: false, context: ['name=world'] });
+
+    const childEnvPath = path.join(traceDir, '0001.includeChild.env');
+    expect(fs.existsSync(childEnvPath)).toBe(true);
+    const envContent = fs.readFileSync(childEnvPath, 'utf8').trim();
+    expect(envContent).toBe(`SOURCE_PATH=${childPath}`);
+    const childMarkup = fs.readFileSync(path.join(traceDir, '0001.includeChild.poml'), 'utf8').trim();
+    expect(childMarkup).toBe('<p>hello {{name}}</p>');
+    const childContext = JSON.parse(fs.readFileSync(path.join(traceDir, '0001.includeChild.context.json'), 'utf8'));
+    expect(childContext).toEqual({ name: 'world' });
+    expect(fs.existsSync(path.join(traceDir, '0001.includeChild.result.json'))).toBe(false);
+
+    fs.rmSync(origDir, { recursive: true, force: true });
+
+    const tracedMainPath = path.join(traceDir, '0001.main.poml');
+    const tracedMarkup = fs.readFileSync(tracedMainPath, 'utf8');
+    const rerenderIr = await read(tracedMarkup, undefined, { name: 'world' }, undefined, tracedMainPath);
+    const rerender = write(rerenderIr);
+    expect(rerender).toBe('hello world');
+  });
+
   test('nextIndex skips index when any file with that index exists', async () => {
     // Create a file with index 0001 but different name to test case 1 logic
     fs.writeFileSync(path.join(traceDir, '0001.different.poml'), 'existing file');

packages/poml/util/trace.ts

@@ -104,11 +104,11 @@ export function dumpTrace(
   result?: any,
   sourcePath?: string,
   prettyResult?: string,
-) {
+): string | undefined {
   if (!isTracing()) {
-    return;
+    return undefined;
   }
-  const [_idx, prefix, fd] = nextIndex(sourcePath);
+  const [idx, prefix, fd] = nextIndex(sourcePath);
   try {
     writeSync(fd, markup);
   } finally {
@@ -136,6 +136,30 @@ export function dumpTrace(
       writeFileSync(`${prefix}.result.txt`, prettyResult);
     }
   }
+
+  const basePrefix = path.join(traceDir!, idx.toString().padStart(4, '0'));
+  return basePrefix;
+}
+
+export function dumpTraceInclude(basePrefix: string, markup: string, context?: any, sourcePath?: string) {
+  if (!isTracing()) {
+    return;
+  }
+  const fileName = sourcePath ? path.basename(sourcePath, '.poml') : 'include';
+  const prefix = `${basePrefix}.${fileName}`;
+  writeFileSync(`${prefix}.poml`, markup);
+  if (sourcePath) {
+    writeFileSync(`${prefix}.env`, `SOURCE_PATH=${sourcePath}\n`);
+    const linkPath = `${prefix}.source.poml`;
+    try {
+      symlinkSync(sourcePath, linkPath);
+    } catch {
+      console.warn(`Failed to create symlink for source path: ${sourcePath}`);
+    }
+  }
+  if (context && Object.keys(context).length > 0) {
+    writeFileSync(`${prefix}.context.json`, JSON.stringify(replaceBuffers(context), null, 2));
+  }
 }
 
 if (process.env.POML_TRACE) {