From e80f882c7592f2a4981748ebedc9fe416a333da3 Mon Sep 17 00:00:00 2001 From: Yuge Zhang Date: Wed, 3 Sep 2025 22:10:58 -0700 Subject: [PATCH] trace included files --- packages/poml/file.tsx | 53 +++++++++++++++++++++++++++--- packages/poml/index.ts | 21 ++++++------ packages/poml/tests/trace.test.tsx | 28 ++++++++++++++++ packages/poml/util/trace.ts | 30 +++++++++++++++-- 4 files changed, 114 insertions(+), 18 deletions(-) diff --git a/packages/poml/file.tsx b/packages/poml/file.tsx index e01e4909..d4fe90f3 100644 --- a/packages/poml/file.tsx +++ b/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 <>; diff --git a/packages/poml/index.ts b/packages/poml/index.ts index c02d92e4..c34514eb 100644 --- a/packages/poml/index.ts +++ b/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 })); } diff --git a/packages/poml/tests/trace.test.tsx b/packages/poml/tests/trace.test.tsx index a1eaaef8..aba8aba7 100644 --- a/packages/poml/tests/trace.test.tsx +++ b/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, ''); + + 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('

hello {{name}}

'); + 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'); diff --git a/packages/poml/util/trace.ts b/packages/poml/util/trace.ts index cf30c458..5c981623 100644 --- a/packages/poml/util/trace.ts +++ b/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) {