microsoft · ultmaster · Jul 25, 2025 · Jul 25, 2025 · Jul 25, 2025 · Jul 25, 2025
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -40,7 +40,7 @@ jobs:
       if: runner.os == 'Linux'
     - run: npm test
       if: runner.os != 'Linux'
-    - run: python -m pytest python/tests
+    - run: python -m pytest -v python/tests
     - run: xvfb-run -a npm run compile && xvfb-run -a npm run test-vscode
       if: runner.os == 'Linux'
     - run: npm run compile && npm run test-vscode

diff --git a/packages/poml-vscode/panel/panel.ts b/packages/poml-vscode/panel/panel.ts
@@ -1,5 +1,6 @@
 import * as vscode from 'vscode';
 import * as path from 'path';
+import * as fs from 'fs';
 
 import { WebviewConfig, WebviewState, WebviewUserOptions, WebviewMessage, PreviewMethodName, PreviewParams, PreviewResponse } from './types';
 import { headlessPomlVscodePanelContent, pomlVscodePanelContent } from './content';
@@ -126,10 +127,6 @@ export class POMLWebviewPanel {
     this._pomlUri = resource;
     this._locked = locked;
     this._userOptions = userOptions;
-    this.resourceOptions.set(this._pomlUri.fsPath, {
-      contexts: [...(userOptions.contexts ?? [])],
-      stylesheets: [...(userOptions.stylesheets ?? [])],
-    });
     this.editor = webview;
 
     this.editor.onDidDispose(
@@ -259,13 +256,9 @@ export class POMLWebviewPanel {
     if (isResourceChange) {
       clearTimeout(this.throttleTimer);
       this.throttleTimer = undefined;
-    }
 
-    if (isResourceChange) {
-      this.resourceOptions.set(this._pomlUri.fsPath, {
-        contexts: [...(this._userOptions.contexts ?? [])],
-        stylesheets: [...(this._userOptions.stylesheets ?? [])],
-      });
+      // Do not save resourceOptions for this._pomlUri
+      // because it may have never changed.
       const saved = this.resourceOptions.get(resource.fsPath);
       if (saved) {
         this._userOptions.contexts = [...saved.contexts];
@@ -277,6 +270,7 @@ export class POMLWebviewPanel {
     }
 
     this._pomlUri = resource;
+    this.autoAddAssociatedFiles(resource.fsPath);
 
     // Schedule update if none is pending
     if (!this.throttleTimer) {
@@ -590,4 +584,35 @@ export class POMLWebviewPanel {
     this.editor.webview.postMessage({ type: WebviewMessage.UpdateUserOptions, options: this._userOptions, source: resource.toString() });
   }
 
+  private autoAddAssociatedFiles(resourcePath: string) {
+    if (this.resourceOptions.has(resourcePath)) {
+      // If we already have options for this resource, no need to add again.
+      return false;
+    }
+
+    const add = (arr: string[], file: string) => {
+      if (fs.existsSync(file) && !arr.includes(file)) {
+        arr.push(file);
+        return true;
+      }
+      return false;
+    };
+
+    let changed = false;
+
+    if (resourcePath.endsWith('.poml')) {
+      const base = resourcePath.replace(/(\.source)?\.poml$/i, '');
+      changed = add(this._userOptions.contexts, `${base}.context.json`) || changed;
+      changed = add(this._userOptions.stylesheets, `${base}.stylesheet.json`) || changed;
+    }
+
+    if (changed) {
+      this.resourceOptions.set(resourcePath, {
+        contexts: [...(this._userOptions.contexts ?? [])],
+        stylesheets: [...(this._userOptions.stylesheets ?? [])],
+      });
+    }
+    return changed;
+  }
+
 }
diff --git a/packages/poml/file.tsx b/packages/poml/file.tsx
@@ -18,6 +18,7 @@ import {
 import { AnyValue, deepMerge, parseText, readSource } from './util';
 import { StyleSheetProvider, ErrorCollection } from './base';
 import { getSuggestions } from './util/xmlContentAssist';
+import { existsSync, readFileSync } from 'fs';
 import path from 'path';
 
 export interface PomlReaderOptions {
@@ -65,6 +66,21 @@ export class PomlFile {
     };
     this.text = this.config.crlfToLf ? text.replace(/\r\n/g, '\n') : text;
     this.sourcePath = sourcePath;
+    if (this.sourcePath) {
+      const envFile = this.sourcePath.replace(/(source\.)?\.poml$/i, '.env');
+      if (existsSync(envFile)) {
+        try {
+          const envText = readFileSync(envFile, 'utf8');
+          const match = envText.match(/^SOURCE_PATH=(.*)$/m);
+          if (match) {
+            // The real source path is specified in the .env file.
+            this.sourcePath = match[1];
+          }
+        } catch {
+          /* ignore */
+        }
+      }
+    }
 
     this.documentRange = { start: 0, end: text.length - 1 };
     let { ast, cst, tokenVector, errors } = this.readXml(text);

diff --git a/packages/poml/index.ts b/packages/poml/index.ts
@@ -187,7 +187,7 @@ export async function commandLine(args: CliArgs) {
 
   if (isTracing()) {
     try {
-      dumpTrace(input, context, stylesheet, result);
+      dumpTrace(input, context, stylesheet, result, sourcePath);
     } catch (err: any) {
       ErrorCollection.add(new SystemError('Failed to dump trace', { cause: err }));
     }

diff --git a/packages/poml/tests/trace.test.tsx b/packages/poml/tests/trace.test.tsx
@@ -2,7 +2,7 @@ import * as fs from 'fs';
 import * as os from 'os';
 import * as path from 'path';
 import { describe, beforeEach, afterEach, test, expect } from '@jest/globals';
-import { commandLine, setTrace, clearTrace, parseJsonWithBuffers, dumpTrace } from 'poml';
+import { commandLine, setTrace, clearTrace, parseJsonWithBuffers, dumpTrace, read, write } from 'poml';
 
 function stringifyWithBuffers(obj: any): string {
   return JSON.stringify(obj, (_k, v) => {
@@ -27,15 +27,36 @@ describe('trace dumps', () => {
   test('unused buffer in context is dumped', async () => {
     const buffer = fs.readFileSync(path.join(__dirname, 'assets', 'tomCat.jpg'));
     dumpTrace('<p></p>', { img: buffer });
-    const raw = fs.readFileSync(path.join(traceDir, '0001_context.json'), 'utf8');
+    const raw = fs.readFileSync(path.join(traceDir, '0001.context.json'), 'utf8');
     expect(raw).toContain('__base64__');
   });
 
   test('document result includes base64', async () => {
     const markup = `<Document src="${path.join(__dirname, 'assets', 'sampleWord.docx')}" />`;
     await commandLine({ input: markup, speakerMode: false });
-    const result = parseJsonWithBuffers(fs.readFileSync(path.join(traceDir, '0001_result.json'), 'utf8'));
+    const result = parseJsonWithBuffers(fs.readFileSync(path.join(traceDir, '0001.result.json'), 'utf8'));
     const images = JSON.stringify(result).includes('base64');
     expect(images).toBe(true);
   });
+
+  test('env file records source path and enables include', async () => {
+    const origDir = fs.mkdtempSync(path.join(os.tmpdir(), 'orig-'));
+    const mainPath = path.join(origDir, 'main.poml');
+    fs.copyFileSync(path.join(__dirname, 'assets', 'includeChild.poml'), path.join(origDir, 'includeChild.poml'));
+    fs.writeFileSync(mainPath, '<poml><include src="includeChild.poml"/></poml>');
+
+    await commandLine({ file: mainPath, speakerMode: false, context: ['name=world'] });
+
+    const envContent = fs.readFileSync(path.join(traceDir, '0001.main.env'), 'utf8').trim();
+    expect(envContent).toBe(`SOURCE_PATH=${mainPath}`);
+
+    const tracedMarkupPath = path.join(traceDir, '0001.main.poml');
+    const traced = fs.readFileSync(tracedMarkupPath, 'utf8');
+    const rerenderIr = await read(traced, undefined, { name: 'world' }, undefined, tracedMarkupPath);
+    const rerender = write(rerenderIr);
+    expect(fs.existsSync(path.join(traceDir, '0001.main.source.poml'))).toBe(true);
+    expect(rerender).toBe('hello world');
+
+    fs.rmSync(origDir, { recursive: true, force: true });
+  });
 });
diff --git a/packages/poml/util/trace.ts b/packages/poml/util/trace.ts
@@ -1,4 +1,4 @@
-import { mkdirSync, writeFileSync, openSync, closeSync, writeSync, readFileSync } from 'fs';
+import { mkdirSync, writeFileSync, openSync, closeSync, writeSync, symlinkSync } from 'fs';
 import path from 'path';
 
 interface Base64Wrapper { __base64__: string }
@@ -60,14 +60,15 @@ export function isTracing(): boolean {
   return traceEnabled && !!traceDir;
 }
 
-function nextIndex(): [number, string, number] {
+function nextIndex(sourcePath?: string): [number, string, number] {
   if (!traceDir) {
     return [0, '', -1];
   }
+  const fileName = sourcePath ? path.basename(sourcePath, '.poml') : '';
   for (let i = 1; ; i++) {
     const idxStr = i.toString().padStart(4, '0');
-    const prefix = path.join(traceDir, idxStr);
-    const filePath = `${prefix}_markup.poml`;
+    const prefix = path.join(traceDir, idxStr) + (fileName ? `.${fileName}` : '');
+    const filePath = `${prefix}.poml`;
     try {
       const fd = openSync(filePath, 'wx');
       return [i, prefix, fd];
@@ -80,24 +81,34 @@ function nextIndex(): [number, string, number] {
   }
 }
 
-export function dumpTrace(markup: string, context?: any, stylesheet?: any, result?: any) {
+export function dumpTrace(markup: string, context?: any, stylesheet?: any, result?: any, sourcePath?: string) {
   if (!isTracing()) {
     return;
   }
-  const [_idx, prefix, fd] = nextIndex();
+  const [_idx, prefix, fd] = nextIndex(sourcePath);
   try {
     writeSync(fd, markup);
   } finally {
     closeSync(fd);
   }
+  if (sourcePath) {
+    const envFile = `${prefix}.env`;
+    writeFileSync(envFile, `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));
+    writeFileSync(`${prefix}.context.json`, JSON.stringify(replaceBuffers(context), null, 2));
   }
   if (stylesheet && Object.keys(stylesheet).length > 0) {
-    writeFileSync(`${prefix}_stylesheet.json`, JSON.stringify(replaceBuffers(stylesheet), null, 2));
+    writeFileSync(`${prefix}.stylesheet.json`, JSON.stringify(replaceBuffers(stylesheet), null, 2));
   }
   if (result !== undefined) {
-    writeFileSync(`${prefix}_result.json`, JSON.stringify(replaceBuffers(result), null, 2));
+    writeFileSync(`${prefix}.result.json`, JSON.stringify(replaceBuffers(result), null, 2));
   }
 }
 

diff --git a/python/tests/test_basic.py b/python/tests/test_basic.py
@@ -99,8 +99,8 @@ def test_trace_directory(tmp_path: Path):
     result = poml("<p>Dir</p>")
     set_trace(False)
     assert result == [{"speaker": "human", "content": "Dir"}]
-    files = list(run_dir.glob("*_markup.poml"))
-    assert len(files) == 1
+    files = list(run_dir.glob("*.poml"))
+    assert len(files) == 2
 
 
 def test_trace_directory_name_format(tmp_path: Path):
@@ -126,7 +126,7 @@ def test_multiprocessing_trace(tmp_path: Path):
         p.join()
     set_trace(False)
     os.environ.pop("POML_TRACE", None)
-    assert len(list(run_dir.glob("*_markup.poml"))) == 3
+    assert len(list(run_dir.glob("*.poml"))) == 6
 
 
 def test_envvar_autotrace(tmp_path: Path):
@@ -135,4 +135,4 @@ def test_envvar_autotrace(tmp_path: Path):
     env["POML_TRACE"] = str(trace_dir)
     script = "from poml import poml; poml('<p>E</p>')"
     subprocess.check_call([sys.executable, "-c", script], env=env)
-    assert any(f.name.endswith("_markup.poml") for f in trace_dir.iterdir())
+    assert any(f.name.endswith(".poml") for f in trace_dir.iterdir())