ai-terminal-assistant/packages/core/tests/unit/agent/prompt-template.test.ts

import { describe, it, expect, beforeEach } from 'vitest';
import {
  renderPromptTemplate,
  createDefaultContext,
  createPlanContext,
  createToolDescriptionContext,
  DEFAULT_TOOL_NAMES,
} from '../../../src/template/index.js';
import type { PromptContext } from '../../../src/template/types.js';

// Alias for backward compatibility in tests
const renderTemplate = renderPromptTemplate;

describe('Prompt Template System', () => {
  let context: PromptContext;

  beforeEach(() => {
    context = createDefaultContext({
      plan: {
        isActive: true,
        planExists: false,
        planFilePath: '/tmp/test-plan.md',
        allowedWritePaths: ['/tmp/*'],
      },
      env: {
        workdir: '/test/project',
        isGitRepo: true,
        platform: 'darwin',
        date: '2024-01-01',
        homeDir: '/home/test',
      },
      agent: {
        name: 'test-agent',
        mode: 'primary',
        isSubagent: false,
      },
    });
  });

  describe('renderTemplate', () => {
    it('should replace simple variables', () => {
      const template = 'Use ${tools.glob} to find files';
      const result = renderTemplate(template, context);
      expect(result).toBe('Use glob to find files');
    });

    it('should replace nested properties', () => {
      const template = 'Plan file: ${plan.planFilePath}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Plan file: /tmp/test-plan.md');
    });

    it('should handle multiple variables in one template', () => {
      const template = 'Use ${tools.glob}, ${tools.grep}, and ${tools.read}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Use glob, grep, and read_file');
    });

    it('should handle ternary expressions - false condition', () => {
      const template = '${plan.planExists ? "exists" : "not exists"}';
      const result = renderTemplate(template, context);
      expect(result).toBe('not exists');
    });

    it('should handle ternary expressions - true condition', () => {
      context.plan.planExists = true;
      const template = '${plan.planExists ? "exists" : "not exists"}';
      const result = renderTemplate(template, context);
      expect(result).toBe('exists');
    });

    it('should handle nested variables in ternary expressions', () => {
      const template =
        '${plan.planExists ? "File at ${plan.planFilePath}" : "Create at ${plan.planFilePath}"}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Create at /tmp/test-plan.md');
    });

    it('should handle boolean values', () => {
      const template = 'Is git repo: ${env.isGitRepo}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Is git repo: true');
    });

    it('should handle undefined variables gracefully', () => {
      const template = 'Value: ${custom.undefined}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Value: ');
    });

    it('should preserve text without variables', () => {
      const template = 'No variables here';
      const result = renderTemplate(template, context);
      expect(result).toBe('No variables here');
    });

    it('should handle complex nested ternary with tool references', () => {
      context.plan.planExists = true;
      const template =
        '${plan.planExists ? "Edit with ${tools.edit}" : "Write with ${tools.write}"}';
      const result = renderTemplate(template, context);
      expect(result).toBe('Edit with edit_file');
    });
  });

  describe('createDefaultContext', () => {
    it('should create context with default tool names', () => {
      const ctx = createDefaultContext();
      expect(ctx.tools.glob).toBe('glob');
      expect(ctx.tools.grep).toBe('grep');
      expect(ctx.tools.read).toBe('read_file');
      expect(ctx.tools.write).toBe('write_file');
      expect(ctx.tools.bash).toBe('bash');
    });

    it('should allow overriding tool names', () => {
      const ctx = createDefaultContext({
        tools: {
          ...DEFAULT_TOOL_NAMES,
          glob: 'custom_glob',
        },
      });
      expect(ctx.tools.glob).toBe('custom_glob');
      expect(ctx.tools.grep).toBe('grep'); // Others unchanged
    });

    it('should set default plan values', () => {
      const ctx = createDefaultContext();
      expect(ctx.plan.isActive).toBe(false);
      expect(ctx.plan.planExists).toBe(false);
    });
  });

  describe('createPlanContext', () => {
    it('should create context with plan mode active', () => {
      const ctx = createPlanContext({});
      expect(ctx.plan.isActive).toBe(true);
      expect(ctx.agent.name).toBe('plan');
    });

    it('should set isSubagent correctly', () => {
      const ctx = createPlanContext({ isSubagent: true });
      expect(ctx.agent.isSubagent).toBe(true);
      expect(ctx.agent.mode).toBe('subagent');
    });

    it('should use custom workdir', () => {
      const ctx = createPlanContext({ workdir: '/custom/path' });
      expect(ctx.env.workdir).toBe('/custom/path');
    });

    it('should use custom plan file path', () => {
      const ctx = createPlanContext({ planFilePath: '/custom/plan.md' });
      expect(ctx.plan.planFilePath).toBe('/custom/plan.md');
    });
  });

  describe('Edge Cases', () => {
    it('should handle empty template', () => {
      const result = renderTemplate('', context);
      expect(result).toBe('');
    });

    it('should handle template with only variable', () => {
      const result = renderTemplate('${tools.bash}', context);
      expect(result).toBe('bash');
    });

    it('should handle unclosed variable syntax', () => {
      const result = renderTemplate('${tools.bash', context);
      expect(result).toBe('${tools.bash');
    });

    it('should handle escaped-like patterns (no real escaping)', () => {
      // Note: This system doesn't support escaping, so $$ or similar won't work
      const result = renderTemplate('Price: $100', context);
      expect(result).toBe('Price: $100');
    });

    it('should handle deeply nested properties', () => {
      const result = renderTemplate('${agent.mode}', context);
      expect(result).toBe('primary');
    });
  });

  describe('createToolDescriptionContext', () => {
    it('should create context with uppercase tool name variables', () => {
      const ctx = createToolDescriptionContext();
      // Check that uppercase tool names are available in custom
      expect(ctx.custom).toBeDefined();
      expect(ctx.custom!['GLOB_TOOL_NAME']).toBe('glob');
      expect(ctx.custom!['GREP_TOOL_NAME']).toBe('grep');
      expect(ctx.custom!['BASH_TOOL_NAME']).toBe('bash');
      expect(ctx.custom!['READ_TOOL_NAME']).toBe('read_file');
      expect(ctx.custom!['WRITE_TOOL_NAME']).toBe('write_file');
      expect(ctx.custom!['EDIT_TOOL_NAME']).toBe('edit_file');
      expect(ctx.custom!['TASK_TOOL_NAME']).toBe('task');
    });

    it('should convert camelCase to UPPER_SNAKE_CASE correctly', () => {
      const ctx = createToolDescriptionContext();
      // askUserQuestion -> ASK_USER_QUESTION_TOOL_NAME
      expect(ctx.custom!['ASK_USER_QUESTION_TOOL_NAME']).toBe('ask_user_question');
      // exitPlanMode -> EXIT_PLAN_MODE_TOOL_NAME
      expect(ctx.custom!['EXIT_PLAN_MODE_TOOL_NAME']).toBe('exit_plan_mode');
      // webSearch -> WEB_SEARCH_TOOL_NAME
      expect(ctx.custom!['WEB_SEARCH_TOOL_NAME']).toBe('web_search');
    });

    it('should support rendering with uppercase tool name variables', () => {
      const ctx = createToolDescriptionContext();
      const template = 'Use ${GREP_TOOL_NAME} for search. Never use ${BASH_TOOL_NAME} grep.';
      const result = renderTemplate(template, ctx);
      expect(result).toBe('Use grep for search. Never use bash grep.');
    });

    it('should support both lowercase and uppercase variable styles', () => {
      const ctx = createToolDescriptionContext();
      const template = 'Tools: ${tools.grep} or ${GREP_TOOL_NAME}';
      const result = renderTemplate(template, ctx);
      expect(result).toBe('Tools: grep or grep');
    });

    it('should preserve overrides', () => {
      const ctx = createToolDescriptionContext({
        agent: {
          name: 'custom-agent',
          mode: 'subagent',
          isSubagent: true,
        },
      });
      expect(ctx.agent.name).toBe('custom-agent');
      expect(ctx.agent.mode).toBe('subagent');
      expect(ctx.agent.isSubagent).toBe(true);
      // Uppercase tool names should still be available
      expect(ctx.custom!['GREP_TOOL_NAME']).toBe('grep');
    });

    it('should render Claude Code style tool description template', () => {
      const ctx = createToolDescriptionContext();
      // Example from Claude Code official prompt
      const template = `A powerful search tool built on ripgrep

Usage:
- ALWAYS use \${GREP_TOOL_NAME} for search tasks. NEVER invoke \`grep\` or \`rg\` as a \${BASH_TOOL_NAME} command.
- Use \${TASK_TOOL_NAME} tool for open-ended searches requiring multiple rounds`;

      const result = renderTemplate(template, ctx);

      expect(result).toContain('ALWAYS use grep for search tasks');
      expect(result).toContain('NEVER invoke `grep` or `rg` as a bash command');
      expect(result).toContain('Use task tool for open-ended searches');
    });
  });
});