feat(core): enhance render-prompt with memoization, slot parsing, and dynamic attributes

packages/core/src/prompts/render-prompt.test.ts

@@ -0,0 +1,454 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, expect, it, vi } from 'vitest';
+import { renderPrompt, p, memoize, parseSlots } from './render-prompt.js';
+import type { PromptContent } from './render-prompt.js';
+
+type TestContext = { name?: string; shouldRender?: boolean };
+
+type TestCase = {
+  desc: string;
+  content: PromptContent<TestContext> | Array<PromptContent<TestContext>>;
+  context: TestContext;
+  contributions?:
+    | Record<string, PromptContent<TestContext>>
+    | Array<Record<string, PromptContent<TestContext>>>;
+  expect: string;
+};
+
+const tests: TestCase[] = [
+  {
+    desc: 'renders a simple string',
+    content: 'Hello world',
+    context: {},
+    expect: 'Hello world',
+  },
+  {
+    desc: 'renders a function content resolving dynamically',
+    content: (ctx) => `Hello ${ctx.name}`,
+    context: { name: 'Alice' },
+    expect: 'Hello Alice',
+  },
+  {
+    desc: 'renders an array of contents with block spacing',
+    content: [['Part 1', 'Part 2']],
+    context: {},
+    expect: 'Part 1\n\nPart 2',
+  },
+  {
+    desc: 'p tag renders inline spacing',
+    content: p`Part ${1} and Part ${2}`,
+    context: {},
+    expect: 'Part 1 and Part 2',
+  },
+  {
+    desc: 'renders a section with heading',
+    content: {
+      heading: 'My Section',
+      content: 'This is the body',
+    },
+    context: {},
+    expect: '# My Section\n\nThis is the body',
+  },
+  {
+    desc: 'renders a section with tag and attrs',
+    content: {
+      tag: 'foo',
+      attrs: { bar: 'baz' },
+      content: 'Hello',
+    },
+    context: {},
+    expect: '<foo bar="baz">\nHello\n</foo>',
+  },
+  {
+    desc: 'conditionally omits rendering a section based on condition',
+    content: [
+      {
+        heading: 'Conditional',
+        condition: (ctx) => ctx.shouldRender ?? false,
+        content: 'This might not appear',
+      },
+      'Always appears',
+    ],
+    context: { shouldRender: false },
+    expect: 'Always appears',
+  },
+  {
+    desc: 'conditionally includes rendering a section based on condition',
+    content: [
+      {
+        heading: 'Conditional',
+        condition: (ctx) => ctx.shouldRender ?? false,
+        content: 'This might not appear',
+      },
+      'Always appears',
+    ],
+    context: { shouldRender: true },
+    expect: '# Conditional\n\nThis might not appear\n\nAlways appears',
+  },
+  {
+    desc: 'conditionally omits rendering a section based on an async condition',
+    content: [
+      {
+        heading: 'Conditional',
+        condition: async (ctx) => ctx.shouldRender ?? false,
+        content: 'This might not appear',
+      },
+      'Always appears',
+    ],
+    context: { shouldRender: false },
+    expect: 'Always appears',
+  },
+  {
+    desc: 'allows dynamic contributions via .contribute() to {slot}',
+    content: [
+      {
+        heading: 'Target Section',
+        content: ['Initial content', { slot: 'target' }],
+      },
+      {
+        heading: 'Other Section',
+        content: 'Other content',
+      },
+    ],
+    contributions: {
+      target: 'Contributed content',
+      missing: 'Should be ignored',
+    },
+    context: {},
+    expect:
+      '# Target Section\n\nInitial content\n\nContributed content\n\n# Other Section\n\nOther content',
+  },
+  {
+    desc: 'handles slot contribution even when missing',
+    content: {
+      heading: 'Target Section',
+      content: ['Initial content', { slot: 'target' }],
+    },
+    context: {},
+    expect: '# Target Section\n\nInitial content',
+  },
+  {
+    desc: 'skips rendering headings and tags if the content is empty',
+    content: [
+      {
+        heading: 'Empty Section',
+        tag: 'empty',
+        content: '',
+      },
+      {
+        heading: 'Empty Array Section',
+        content: [],
+      },
+      {
+        heading: 'Function resolving to empty',
+        content: () => '',
+      },
+      'Visible content',
+    ],
+    context: {},
+    expect: 'Visible content',
+  },
+  {
+    desc: 'handles nested structures properly in contribute',
+    content: {
+      heading: 'Outer',
+      content: [
+        {
+          heading: 'Inner',
+          content: ['Inner content', { slot: 'inner' }],
+        },
+      ],
+    },
+    contributions: {
+      inner: 'Injected into inner',
+    },
+    context: {},
+    expect: '# Outer\n\n## Inner\n\nInner content\n\nInjected into inner',
+  },
+  {
+    desc: 'resolves recursive async functions correctly before filling slots',
+    content: {
+      heading: 'Async Section',
+      content: async () => [
+        'Content from async function',
+        { slot: 'async_slot' },
+      ],
+    },
+    contributions: {
+      async_slot: async () => 'Contributed async content',
+    },
+    context: {},
+    expect:
+      '# Async Section\n\nContent from async function\n\nContributed async content',
+  },
+  {
+    desc: 'collapses 3+ newlines into 2',
+    content: ['First', '\n\n\n\n\n', 'Second'],
+    context: {},
+    expect: 'First\n\nSecond',
+  },
+  {
+    desc: 'does not collapse 3+ newlines inside markdown code fences',
+    content: ['First', '```\n\n\n\n\n```', 'Second'],
+    context: {},
+    expect: 'First\n\n```\n\n\n\n\n```\n\nSecond',
+  },
+  {
+    desc: 'appends multiple contributions to the same slot',
+    content: {
+      heading: 'Multi Section',
+      content: [{ slot: 'multi' }],
+    },
+    contributions: [{ multi: 'First' }, { multi: 'Second' }],
+    context: {},
+    expect: '# Multi Section\n\nFirst\n\nSecond',
+  },
+  {
+    desc: 'appends multiple contributions to an inline slot',
+    content: p`Prefix: ${{ slot: 'inline_multi' }}`,
+    contributions: [{ inline_multi: 'First' }, { inline_multi: 'Second' }],
+    context: {},
+    expect: 'Prefix: FirstSecond',
+  },
+  {
+    desc: 'conditionally omits items when null or undefined are present',
+    content: [
+      'Item 1',
+      null,
+      'Item 2',
+      undefined,
+      {
+        heading: 'Optional Section',
+        content: null,
+      },
+      'Item 3',
+    ],
+    context: {},
+    expect: 'Item 1\n\nItem 2\n\nItem 3',
+  },
+  {
+    desc: 'renders lists with dashes',
+    content: {
+      heading: 'My List',
+      format: 'list',
+      content: ['Apple', 'Banana', null, 'Cherry'],
+    },
+    context: {},
+    expect: '# My List\n\n- Apple\n- Banana\n- Cherry',
+  },
+  {
+    desc: 'supports custom format functions',
+    content: {
+      heading: 'Custom Format',
+      format: (parts) => parts.join(' | '),
+      content: ['A', 'B', undefined, 'C'],
+    },
+    context: {},
+    expect: '# Custom Format\n\nA | B | C',
+  },
+  {
+    desc: 'filters out boolean values',
+    content: ['First', true, 'Second', false, 'Third'],
+    context: {},
+    expect: 'First\n\nSecond\n\nThird',
+  },
+  {
+    desc: 'supports static condition booleans',
+    content: [
+      {
+        heading: 'Shown',
+        condition: true,
+        content: 'This should show',
+      },
+      {
+        heading: 'Hidden',
+        condition: false,
+        content: 'This should hide',
+      },
+    ],
+    context: {},
+    expect: '# Shown\n\nThis should show',
+  },
+  {
+    desc: 'supports lines format',
+    content: {
+      heading: 'Lines List',
+      format: 'lines',
+      content: ['Line 1', 'Line 2', 'Line 3'],
+    },
+    context: {},
+    expect: '# Lines List\n\nLine 1\nLine 2\nLine 3',
+  },
+  {
+    desc: 'renders fallback when content is falsy (boolean)',
+    content: {
+      heading: 'Fallback Boolean',
+      content: false,
+      fallback: 'This is the fallback',
+    },
+    context: {},
+    expect: '# Fallback Boolean\n\nThis is the fallback',
+  },
+  {
+    desc: 'renders fallback when content is falsy (empty string via short-circuit)',
+    content: {
+      heading: 'Fallback Empty String',
+      content: String('') && 'something',
+      fallback: 'Fallback for empty string',
+    },
+    context: {},
+    expect: '# Fallback Empty String\n\nFallback for empty string',
+  },
+  {
+    desc: 'renders fallback when content is an array that filters down to empty',
+    content: {
+      heading: 'Fallback Empty Array',
+      content: [false, null, '', undefined],
+      fallback: 'Fallback for empty array',
+    },
+    context: {},
+    expect: '# Fallback Empty Array\n\nFallback for empty array',
+  },
+  {
+    desc: 'renders normal content and ignores fallback when content is present',
+    content: {
+      heading: 'No Fallback',
+      content: 'Primary content',
+      fallback: 'Should not see this',
+    },
+    context: {},
+    expect: '# No Fallback\n\nPrimary content',
+  },
+  {
+    desc: 'does not render section if both content and fallback are falsy',
+    content: [
+      'Visible start',
+      {
+        heading: 'Double Falsy',
+        content: '',
+        fallback: null,
+      },
+      'Visible end',
+    ],
+    context: {},
+    expect: 'Visible start\n\nVisible end',
+  },
+  {
+    desc: 'renders explicit level overriding depth and resetting children',
+    content: {
+      heading: 'Top Level',
+      content: {
+        heading: 'Nested',
+        level: 4,
+        content: {
+          heading: 'Deep',
+          content: 'Text',
+        },
+      },
+    },
+    context: {},
+    expect: '# Top Level\n\n#### Nested\n\n##### Deep\n\nText',
+  },
+  {
+    desc: 'renders level 0 as # and children as level 1 (#)',
+    content: {
+      heading: 'Level 0',
+      level: 0,
+      content: {
+        heading: 'Level 1',
+        content: 'Text',
+      },
+    },
+    context: {},
+    expect: '# Level 0\n\n# Level 1\n\nText',
+  },
+  {
+    desc: 'resolves dynamic attributes synchronously',
+    content: {
+      tag: 'dynamic',
+      attrs: { static: 'val', dyn: (ctx) => `hello-${ctx.name}` },
+      content: 'Inside',
+    },
+    context: { name: 'Alice' },
+    expect: '<dynamic static="val" dyn="hello-Alice">\nInside\n</dynamic>',
+  },
+  {
+    desc: 'resolves dynamic attributes asynchronously',
+    content: {
+      tag: 'async-dynamic',
+      attrs: { dyn: async (ctx) => `async-${ctx.name}` },
+      content: 'Inside',
+    },
+    context: { name: 'Bob' },
+    expect: '<async-dynamic dyn="async-Bob">\nInside\n</async-dynamic>',
+  },
+];
+
+describe('renderPrompt', () => {
+  it.each(tests)('$desc', async (test) => {
+    const result = await renderPrompt({
+      content: test.content,
+      contributions: test.contributions,
+      context: test.context,
+    });
+    expect(result).toBe(test.expect);
+  });
+});
+
+describe('memoize', () => {
+  it('should cache result per context instance', () => {
+    const resolver = vi.fn((ctx: TestContext) => ctx.name);
+    const memoized = memoize(resolver);
+
+    const ctx1 = { name: 'Alice' };
+    const ctx2 = { name: 'Bob' };
+
+    expect(memoized(ctx1)).toBe('Alice');
+    expect(memoized(ctx1)).toBe('Alice');
+    expect(resolver).toHaveBeenCalledTimes(1);
+
+    expect(memoized(ctx2)).toBe('Bob');
+    expect(memoized(ctx2)).toBe('Bob');
+    expect(resolver).toHaveBeenCalledTimes(2);
+  });
+
+  it('should handle async resolvers', async () => {
+    const resolver = vi.fn(async (ctx: TestContext) => ctx.name);
+    const memoized = memoize(resolver);
+
+    const ctx = { name: 'Async' };
+
+    expect(await memoized(ctx)).toBe('Async');
+    expect(await memoized(ctx)).toBe('Async');
+    expect(resolver).toHaveBeenCalledTimes(1);
+  });
+});
+
+describe('parseSlots', () => {
+  it('should return empty array for empty string', () => {
+    expect(parseSlots('')).toEqual([]);
+  });
+
+  it('should return string array for no slots', () => {
+    expect(parseSlots('Hello World')).toEqual(['Hello World']);
+  });
+
+  it('should parse a single slot', () => {
+    expect(parseSlots('${slot1}')).toEqual([{ slot: 'slot1' }]);
+  });
+
+  it('should parse slots at the start, middle, and end', () => {
+    expect(parseSlots('${first} middle ${second} end ${third}')).toEqual([
+      { slot: 'first' },
+      ' middle ',
+      { slot: 'second' },
+      ' end ',
+      { slot: 'third' },
+    ]);
+  });
+});

packages/core/src/prompts/render-prompt.ts

@@ -0,0 +1,332 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { SystemPromptOptions } from './snippets.js';
+
+export type ContextResolver<C, O> = O | ((ctx: C) => O | Promise<O>);
+
+export type PromptSlot = { slot: string; content?: never };
+
+export type PromptSection<C> = {
+  /** Add a Markdown heading of appropriate level to this section. */
+  heading?: string;
+  /** Explicitly set the markdown heading depth (1 = #, 2 = ##), overriding tree depth.
+   * 0 is valid and will render as # while setting children to level 1. */
+  level?: number;
+  /** If supplied, wrap this section in an XML tag. */
+  tag?: string;
+  /** If supplied, add attributes to the XML section tag. */
+  attrs?: Record<string, ContextResolver<C, string>>;
+  /** Formatting of the content inside this section. Defaults to 'block'. */
+  format?:
+    | 'inline'
+    | 'block'
+    | 'list'
+    | 'lines'
+    | ((parts: string[]) => string);
+
+  /** Condition that must evaluate to true for the section to be rendered. */
+  condition?: boolean | ((ctx: C) => boolean | Promise<boolean>);
+  content: PromptContent<C>;
+  /** Alternate content to render if the primary content resolves to a falsy value. */
+  fallback?: PromptContent<C>;
+};
+
+// The core recursive type.
+// It wraps your 3 base node shapes (string, section, or array) in the resolver.
+export type PromptContent<C> = ContextResolver<
+  C,
+  | string
+  | number
+  | boolean
+  | null
+  | undefined
+  | PromptSection<C>
+  | PromptSlot
+  | Array<PromptContent<C>>
+>;
+
+type BaseContent = string | StaticSection | PromptSlot | BaseContent[];
+type StaticSection = Omit<
+  PromptSection<unknown>,
+  'condition' | 'content' | 'attrs'
+> & {
+  attrs?: Record<string, string>;
+  content: BaseContent;
+};
+
+// Helper to stringify XML attributes cleanly
+function renderAttributes(attrs?: Record<string, string>): string {
+  if (!attrs) return '';
+  return Object.entries(attrs)
+    .map(([key, value]) => ` ${key}="${value}"`)
+    .join('');
+}
+
+export function p<C = SystemPromptOptions>(
+  strings: TemplateStringsArray,
+  ...values: Array<PromptContent<C>>
+): PromptContent<C> {
+  const content = strings.reduce<Array<PromptContent<C>>>(
+    (acc, str, i) => [...acc, str, values[i] ?? ''],
+    [],
+  );
+  return { format: 'inline', content };
+}
+
+export interface RenderPromptOptions<C> {
+  content: PromptContent<C> | Array<PromptContent<C>>;
+  contributions?:
+    | Record<string, PromptContent<C>>
+    | Array<Record<string, PromptContent<C>>>;
+  context: C;
+  options?: { depth?: number };
+}
+
+function normalizeResult(rawResult: string): string {
+  // Normalize newlines: collapse 3+ consecutive newlines into exactly 2
+  // but skip content inside markdown code fences (```)
+  const segments = rawResult.split(/(```[\s\S]*?```)/);
+  return segments
+    .map((segment, index) => {
+      // Even indices are outside code fences, odd indices are inside
+      if (index % 2 === 0) {
+        return segment.replace(/\n{3,}/g, '\n\n');
+      }
+      return segment;
+    })
+    .join('');
+}
+
+const formatBasic = (
+  c: BaseContent | null,
+  depth: number,
+  format: 'inline' | 'block' | 'list' | 'lines' | ((parts: string[]) => string),
+  resolvedContributions: Record<string, BaseContent[]>,
+): string => {
+  if (c === null) return '';
+  if (typeof c === 'string') return c;
+  if (Array.isArray(c)) {
+    const parts = c
+      .map((item) => formatBasic(item, depth, format, resolvedContributions))
+      .filter((p) => p !== '');
+    if (typeof format === 'function') {
+      return format(parts);
+    }
+    if (format === 'list') {
+      return parts.map((p) => '- ' + p).join('\n');
+    }
+    if (format === 'lines') {
+      return parts.join('\n');
+    }
+    return parts.join(format === 'inline' ? '' : '\n\n');
+  }
+  if ('slot' in c) {
+    const slotContributions = resolvedContributions[c.slot];
+    if (!slotContributions || slotContributions.length === 0) return '';
+    return formatBasic(slotContributions, depth, format, resolvedContributions);
+  }
+
+  const section = c;
+  const currentDepth = section.level ?? depth;
+  const sectionFormat = section.format || 'block';
+  const innerContent = formatBasic(
+    section.content,
+    currentDepth + 1,
+    sectionFormat,
+    resolvedContributions,
+  ).trim();
+  if (!innerContent) return '';
+
+  let result = innerContent;
+
+  if (section.tag) {
+    const attrs = renderAttributes(section.attrs);
+    result = `\n<${section.tag}${attrs}>\n${result}\n</${section.tag}>\n`;
+  }
+
+  if (section.heading) {
+    const headingLevel = Math.max(1, Math.min(currentDepth, 6));
+    result = `\n\n${'#'.repeat(headingLevel)} ${section.heading}\n\n${result.trim()}`;
+  }
+
+  return result;
+};
+
+export async function renderPrompt<C = SystemPromptOptions>({
+  content,
+  contributions,
+  context,
+  options,
+}: RenderPromptOptions<C>): Promise<string> {
+  const contents = Array.isArray(content) ? content : [content];
+  const _contributions: Record<string, Array<PromptContent<C>>> = {};
+
+  if (contributions) {
+    const batches = Array.isArray(contributions)
+      ? contributions
+      : [contributions];
+    for (const batch of batches) {
+      for (const [slot, c] of Object.entries(batch)) {
+        _contributions[slot] = _contributions[slot] || [];
+        _contributions[slot].push(c);
+      }
+    }
+  }
+
+  const resolveToBasic = async (
+    c: PromptContent<C>,
+  ): Promise<BaseContent | null> => {
+    if (c === undefined || c === null) return null;
+    if (typeof c === 'function') {
+      const resolved = await c(context);
+      return resolveToBasic(resolved);
+    }
+    if (typeof c === 'boolean') {
+      return null;
+    }
+    if (typeof c === 'string' || typeof c === 'number') {
+      const val = String(c);
+      return val === '' ? null : val;
+    }
+    if (Array.isArray(c)) {
+      const resolved = await Promise.all(c.map((item) => resolveToBasic(item)));
+      const filtered = resolved.filter(
+        (item): item is BaseContent => item !== null,
+      );
+      if (filtered.length === 0) return null;
+      return filtered;
+    }
+    if (typeof c === 'object' && c !== null) {
+      if ('slot' in c) {
+        return c;
+      }
+
+      const section = c;
+      if (section.condition !== undefined) {
+        const shouldRender =
+          typeof section.condition === 'function'
+            ? await section.condition(context)
+            : section.condition;
+        if (!shouldRender) return null;
+      }
+      let resolvedInner = await resolveToBasic(section.content);
+
+      if (
+        resolvedInner === null ||
+        resolvedInner === '' ||
+        (Array.isArray(resolvedInner) && resolvedInner.length === 0)
+      ) {
+        if (section.fallback !== undefined) {
+          resolvedInner = await resolveToBasic(section.fallback);
+        }
+      }
+
+      if (
+        resolvedInner === null ||
+        resolvedInner === '' ||
+        (Array.isArray(resolvedInner) && resolvedInner.length === 0)
+      ) {
+        return null;
+      }
+      let resolvedAttrs: Record<string, string> | undefined = undefined;
+      if (section.attrs) {
+        resolvedAttrs = {};
+        for (const [key, value] of Object.entries(section.attrs)) {
+          resolvedAttrs[key] =
+            typeof value === 'function' ? await value(context) : value;
+        }
+      }
+
+      return {
+        heading: section.heading,
+        level: section.level,
+        tag: section.tag,
+        attrs: resolvedAttrs,
+        format: section.format,
+        content: resolvedInner,
+      };
+    }
+    return null;
+  };
+
+  const resolvedContents = await Promise.all(
+    contents.map((c) => resolveToBasic(c)),
+  );
+
+  const resolvedContributions: Record<string, BaseContent[]> = {};
+  for (const [slot, slotContributions] of Object.entries(_contributions)) {
+    const resolved = await Promise.all(
+      slotContributions.map((c) => resolveToBasic(c)),
+    );
+    resolvedContributions[slot] = resolved.filter(
+      (c): c is BaseContent => c !== null,
+    );
+  }
+
+  const parts = resolvedContents
+    .map((c) =>
+      formatBasic(c, options?.depth ?? 1, 'block', resolvedContributions),
+    )
+    .filter((p) => p !== null && p !== '');
+
+  const rawResult = parts.join('\n\n').trim();
+  return normalizeResult(rawResult);
+}
+
+export function prompt<C = SystemPromptOptions>(
+  ...content: Array<PromptContent<C>>
+): PromptContent<C> {
+  return content.length === 1 ? content[0] : content;
+}
+
+type Resolver<C, T> = (ctx: C) => T | Promise<T>;
+
+/**
+ * Creates a memoized selector that caches its result per context instance.
+ * Ideal for efficiently sharing derived state across a prompt tree.
+ */
+export function memoize<C extends object, T>(
+  resolver: Resolver<C, T>,
+): (ctx: C) => T | Promise<T> {
+  const cache = new WeakMap<C, T | Promise<T>>();
+
+  return (ctx: C) => {
+    if (cache.has(ctx)) {
+      return cache.get(ctx)!;
+    }
+    const result = resolver(ctx);
+    cache.set(ctx, result);
+    return result;
+  };
+}
+
+/**
+ * Parses a string containing placeholders like `${slotName}` into a PromptContent array.
+ * Interleaves literal string segments with `{ slot: 'slotName' }` objects.
+ */
+export function parseSlots<C>(template: string): Array<PromptContent<C>> {
+  if (!template) return [];
+
+  const regex = /\$\{([^}]+)\}/g;
+  const parts: Array<PromptContent<C>> = [];
+  let lastIndex = 0;
+  let match;
+
+  while ((match = regex.exec(template)) !== null) {
+    if (match.index > lastIndex) {
+      parts.push(template.slice(lastIndex, match.index));
+    }
+    parts.push({ slot: match[1] });
+    lastIndex = regex.lastIndex;
+  }
+
+  if (lastIndex < template.length) {
+    parts.push(template.slice(lastIndex));
+  }
+
+  return parts;
+}