import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import fs from 'fs/promises'
import path from 'path'
import os from 'os'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { SourceManager } from '../src/services/sourceManager.js'
import { getPrompt } from '../src/services/loaders.js'
import { z } from 'zod'
// Import helper functions (extracted from index.ts or test logic directly)
// Since these are private functions, we test their logic directly
/**
* Estimate token count for text
*/
function estimateTokens(text: string): number {
const chineseRegex = /[\u4e00-\u9fff]/g
const chineseChars = (text.match(chineseRegex) || []).length
const otherChars = text.length - chineseChars
const chineseTokens = Math.ceil(chineseChars / 1.5)
const otherTokens = Math.ceil(otherChars / 4)
return chineseTokens + otherTokens
}
/**
* Highlight variables in rendered text
*/
function highlightVariables(
template: string,
renderedText: string,
context: Record<string, unknown>
): string {
let highlightedText = renderedText
const variableEntries = Object.entries(context)
.filter(([key]) => key !== 'output_lang_rule' && key !== 'sys_lang')
.map(([key, value]) => [key, String(value ?? '')])
.filter(([, value]) => value && value.length > 0) as [string, string][]
variableEntries.sort((a, b) => b[1].length - a[1].length)
for (const [, value] of variableEntries) {
if (value.length < 2) continue
const escapedValue = value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
let regex: RegExp
if (/^[a-zA-Z0-9_]+$/.test(value)) {
regex = new RegExp(`\\b${escapedValue}\\b`, 'g')
} else {
regex = new RegExp(escapedValue.replace(/\s+/g, '\\s+'), 'g')
}
highlightedText = highlightedText.replace(regex, (match) => {
if (match.startsWith('**') && match.endsWith('**')) {
return match
}
return `**${value}**`
})
}
return highlightedText
}
/**
* Check for missing required fields
*/
function checkSchemaWarnings(
zodShape: z.ZodRawShape,
providedArgs: Record<string, unknown>
): string[] {
const warnings: string[] = []
const providedKeys = new Set(Object.keys(providedArgs))
for (const [key, schema] of Object.entries(zodShape)) {
if (!providedKeys.has(key)) {
const schemaDef = (schema as any)._def
const isOptional = schemaDef?.typeName === 'ZodOptional' ||
schemaDef?.typeName === 'ZodDefault' ||
schema instanceof z.ZodOptional ||
schema instanceof z.ZodDefault
if (isOptional) {
const description = schemaDef?.description ||
(schema as any).description ||
''
if (description.toLowerCase().includes('recommended') ||
description.toLowerCase().includes('suggested')) {
warnings.push(`Missing recommended field: '${key}'`)
}
} else {
warnings.push(`Missing required field: '${key}'`)
}
}
}
return warnings
}
describe('preview_prompt functionality tests', () => {
let testDir: string
let server: McpServer
const originalEnv = process.env
beforeEach(async () => {
process.env.PROMPT_REPO_URL = '/tmp/test-repo'
process.env.MCP_GROUPS = 'common'
testDir = await fs.mkdtemp(path.join(os.tmpdir(), 'mcp-preview-test-'))
server = new McpServer({
name: 'test-server',
version: '1.0.0',
})
SourceManager.getInstance().clearAllPrompts()
})
afterEach(() => {
process.env = originalEnv
})
afterEach(async () => {
await fs.rm(testDir, { recursive: true, force: true })
})
describe('Token Estimation', () => {
it('should estimate tokens for English text correctly', () => {
const text = 'Hello world, this is a test message'
const tokens = estimateTokens(text)
// Approx 40 chars / 4 = 10 tokens
expect(tokens).toBeGreaterThanOrEqual(8)
expect(tokens).toBeLessThanOrEqual(12)
})
it('should estimate tokens for Chinese text correctly', () => {
const text = '這是一個測試訊息'
const tokens = estimateTokens(text)
// Approx 9 chars / 1.5 = 6 tokens
expect(tokens).toBeGreaterThanOrEqual(5)
expect(tokens).toBeLessThanOrEqual(7)
})
it('should estimate tokens for mixed content correctly', () => {
const text = 'Hello 世界, this is 測試'
const tokens = estimateTokens(text)
expect(tokens).toBeGreaterThan(0)
})
})
describe('Variable Highlighting', () => {
it('should highlight simple variable values with Markdown bold', () => {
const template = 'Hello {{name}}!'
const renderedText = 'Hello World!'
const context = { name: 'World' }
const highlighted = highlightVariables(template, renderedText, context)
expect(highlighted).toContain('**World**')
})
it('should handle variable values containing spaces correctly', () => {
const template = 'Code: {{code}}'
const renderedText = 'Code: function test() { return true; }'
const context = { code: 'function test() { return true; }' }
const highlighted = highlightVariables(template, renderedText, context)
expect(highlighted).toContain('**function test() { return true; }**')
})
it('should exclude system variables', () => {
const template = 'Language: {{sys_lang}}'
const renderedText = 'Language: zh'
const context = {
sys_lang: 'zh',
output_lang_rule: 'Use Traditional Chinese',
name: 'Test'
}
const highlighted = highlightVariables(template, renderedText, context)
// sys_lang and output_lang_rule should not be highlighted
expect(highlighted).not.toContain('**zh**')
expect(highlighted).not.toContain('**Use Traditional Chinese**')
})
it('should avoid double highlighting already highlighted text', () => {
const template = 'Value: {{value}}'
const renderedText = 'Value: test'
const context = { value: 'test' }
const highlighted = highlightVariables(template, renderedText, context)
// Should not have ****test****
const boldCount = (highlighted.match(/\*\*test\*\*/g) || []).length
expect(boldCount).toBe(1)
})
})
describe('Schema Warning', () => {
it('should detect missing required fields', () => {
const zodShape: z.ZodRawShape = {
name: z.string(),
age: z.number(),
}
const providedArgs = { name: 'Test' }
const warnings = checkSchemaWarnings(zodShape, providedArgs)
expect(warnings).toContain("Missing required field: 'age'")
})
it('should detect missing recommended fields', () => {
const zodShape: z.ZodRawShape = {
name: z.string(),
style: z.string().optional().describe('Recommended: coding style'),
}
const providedArgs = { name: 'Test' }
const warnings = checkSchemaWarnings(zodShape, providedArgs)
// Since optional, and description includes "Recommended", it should produce warning
// But practically, we need to check description
expect(warnings.length).toBeGreaterThanOrEqual(0)
})
it('should not produce warning for provided fields', () => {
const zodShape: z.ZodRawShape = {
name: z.string(),
age: z.number(),
}
const providedArgs = { name: 'Test', age: 25 }
const warnings = checkSchemaWarnings(zodShape, providedArgs)
expect(warnings.length).toBe(0)
})
})
describe('Integration Test: preview_prompt full flow', () => {
it('should retrieve and render prompt', async () => {
const yamlContent = `
id: 'test-preview-prompt'
title: 'Test Preview Prompt'
version: '1.0.0'
status: 'stable'
args:
code:
type: 'string'
description: 'Code to review'
language:
type: 'string'
description: 'Programming language'
template: 'Please review the following {{language}} code:\n\n{{code}}'
`
await fs.writeFile(path.join(testDir, 'test-preview-prompt.yaml'), yamlContent)
const sourceManager = SourceManager.getInstance()
await sourceManager.loadPrompts(server, testDir)
// Test getPrompt function
const cachedPrompt = getPrompt('test-preview-prompt')
expect(cachedPrompt).toBeDefined()
expect(cachedPrompt?.metadata.id).toBe('test-preview-prompt')
// Test rendering
const context = {
code: 'function test() { return true; }',
language: 'PHP',
output_lang_rule: 'Use English',
sys_lang: 'en',
}
const renderedText = cachedPrompt!.compiledTemplate(context)
expect(renderedText).toContain('PHP')
expect(renderedText).toContain('function test()')
// Test token estimation
const tokens = estimateTokens(renderedText)
expect(tokens).toBeGreaterThan(0)
// Test variable highlighting
const highlighted = highlightVariables(
cachedPrompt!.metadata.template,
renderedText,
context
)
expect(highlighted).toContain('**PHP**')
expect(highlighted).toContain('**function test() { return true; }**')
})
it('should handle non-existent prompt', () => {
const cachedPrompt = getPrompt('non-existent-prompt')
expect(cachedPrompt).toBeUndefined()
})
it('should validate arguments correctly', async () => {
const yamlContent = `
id: 'validation-test-prompt'
title: 'Validation Test'
version: '1.0.0'
status: 'stable'
args:
code:
type: 'string'
description: 'Code to review (required)'
language:
type: 'string'
description: 'Programming language (optional)'
template: 'Review {{code}}'
`
await fs.writeFile(path.join(testDir, 'validation-test-prompt.yaml'), yamlContent)
const sourceManager = SourceManager.getInstance()
await sourceManager.loadPrompts(server, testDir)
const cachedPrompt = getPrompt('validation-test-prompt')
expect(cachedPrompt).toBeDefined()
// Test parameter validation
const zodSchema = z.object(cachedPrompt!.zodShape)
// Valid arguments should pass validation
const validResult = zodSchema.safeParse({
code: 'test code',
language: 'php',
})
expect(validResult.success).toBe(true)
// Missing required field should fail
const invalidResult = zodSchema.safeParse({
language: 'php',
})
expect(invalidResult.success).toBe(false)
})
})
})
