1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
use super::defaults;
use super::metadata::PromptTemplate;
use super::template::{parse_template, render_template};
use anyhow::{Context, Result};
use kodegen_config::KodegenConfig;
use kodegen_mcp_tool::error::McpError;
use log::{debug, info, warn};
use std::collections::HashMap;
use std::io::ErrorKind;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::SystemTime;
use tokio::fs;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::sync::RwLock;
/// Cached template with file modification time for validation
struct CachedTemplate {
template: PromptTemplate,
file_mtime: SystemTime,
}
#[derive(Clone)]
pub struct PromptManager {
prompts_dir: PathBuf,
cache: Arc<RwLock<HashMap<String, CachedTemplate>>>,
}
impl Default for PromptManager {
fn default() -> Self {
Self::new()
}
}
impl PromptManager {
/// Create new prompt manager (synchronous constructor)
#[must_use]
pub fn new() -> Self {
let prompts_dir =
get_prompts_directory().unwrap_or_else(|_| PathBuf::from(".kodegen/prompts"));
Self {
prompts_dir,
cache: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Initialize the prompt manager (async initialization)
///
/// Call this after `new()` to perform async setup operations.
pub async fn init(&self) -> Result<(), McpError> {
// Ensure directory exists (async)
fs::create_dir_all(&self.prompts_dir)
.await
.with_context(|| {
format!(
"Failed to create prompts directory: {}",
self.prompts_dir.display()
)
})
.map_err(McpError::Other)?;
// Initialize default prompts if directory is empty (async)
if let Err(e) = initialize_default_prompts(&self.prompts_dir).await {
warn!("Failed to initialize default prompts: {e}");
// Don't fail - user can add prompts manually
}
Ok(())
}
/// List all available prompts (async)
pub async fn list_prompts(&self) -> Result<Vec<PromptTemplate>> {
let mut prompts = Vec::new();
let mut entries = fs::read_dir(&self.prompts_dir).await.with_context(|| {
format!(
"Failed to read prompts directory: {}",
self.prompts_dir.display()
)
})?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
// CHANGE 1: Check file type first (reject symlinks and directories)
let file_type = match entry.file_type().await {
Ok(ft) => ft,
Err(e) => {
warn!("Failed to get file type for {}: {e}", path.display());
continue;
}
};
// CHANGE 2: Skip non-regular files (directories, symlinks, etc.)
if !file_type.is_file() {
debug!("Skipping non-file entry: {}", path.display());
continue;
}
// CHANGE 3: Check for .j2.md extension (not just .md)
let filename_str = match path.file_name().and_then(|s| s.to_str()) {
Some(name) if name.ends_with(".j2.md") => name,
_ => continue, // Skip files that don't match pattern
};
// CHANGE 4: Extract stem by removing ".j2.md" suffix (6 chars)
let stem = &filename_str[..filename_str.len() - 6];
// Validate prompt name before attempting load (reuses existing validation)
if !is_valid_prompt_name(stem) {
warn!("Invalid prompt filename (skipping): {stem}");
continue;
}
// Load prompt (now guaranteed to be safe, regular file)
match self.load_prompt(stem).await {
Ok(template) => prompts.push(template),
Err(e) => {
warn!("Failed to load prompt '{stem}': {e}");
}
}
}
Ok(prompts)
}
/// Load a specific prompt by filename (async)
pub async fn load_prompt(&self, name: &str) -> Result<PromptTemplate> {
// Validate name to prevent path traversal
validate_prompt_name(name)?;
let path = self.prompts_dir.join(format!("{name}.j2.md"));
// Step 1: Check cache with read lock (allows concurrent reads)
{
let cache = self.cache.read().await;
if let Some(cached) = cache.get(name) {
// Verify file hasn't been modified since caching
if let Ok(current_meta) = fs::metadata(&path).await
&& let Ok(current_mtime) = current_meta.modified()
&& current_mtime == cached.file_mtime {
// Cache hit: file unchanged, return cached template
return Ok(cached.template.clone());
}
// Cache stale: file modified, fall through to reload
}
// Cache miss: template not cached, fall through to load
} // Read lock dropped here
// Step 2: Cache miss or stale - load from disk
let content = fs::read_to_string(&path)
.await
.with_context(|| format!("Failed to read prompt: {name}"))?;
let metadata = fs::metadata(&path).await?;
let file_mtime = metadata.modified()?;
let template = parse_template(name, &content)?;
// Step 3: Update cache with write lock
{
let mut cache = self.cache.write().await;
cache.insert(
name.to_string(),
CachedTemplate {
template: template.clone(),
file_mtime,
},
);
} // Write lock dropped here
Ok(template)
}
/// Save a new prompt (async)
pub async fn add_prompt(&self, name: &str, content: &str) -> Result<()> {
// Validate name (prevent path traversal)
validate_prompt_name(name)?;
// Validate content syntax
super::validation::validate_prompt_file(content)?;
let path = self.prompts_dir.join(format!("{name}.j2.md"));
// Atomic create-new operation - fails if file already exists
match OpenOptions::new()
.write(true)
.create_new(true) // Atomic: fails if file exists
.open(&path)
.await
{
Ok(mut file) => {
// File created successfully, write content
file.write_all(content.as_bytes())
.await
.with_context(|| format!("Failed to write prompt: {name}"))?;
file.flush()
.await
.with_context(|| format!("Failed to flush prompt: {name}"))?;
// Sync to disk for durability (survive power loss)
file.sync_all()
.await
.with_context(|| format!("Failed to sync prompt to disk: {name}"))?;
// Invalidate cache after successful write
self.invalidate_cache(name).await;
Ok(())
}
Err(e) if e.kind() == ErrorKind::AlreadyExists => {
// File already exists - provide friendly error message
anyhow::bail!("Prompt '{name}' already exists. Use edit_prompt to modify.")
}
Err(e) => {
// Other IO error (permissions, disk full, etc.)
Err(e).with_context(|| format!("Failed to create prompt: {name}"))?
}
}
}
/// Update an existing prompt (async)
pub async fn edit_prompt(&self, name: &str, content: &str) -> Result<()> {
validate_prompt_name(name)?;
super::validation::validate_prompt_file(content)?;
let path = self.prompts_dir.join(format!("{name}.j2.md"));
// Atomic update-only operation - fails if file doesn't exist
match OpenOptions::new()
.write(true)
.truncate(true)
.create(false) // CRITICAL: Fail if file doesn't exist (edit-only semantics)
.open(&path)
.await
{
Ok(mut file) => {
// Write new content to existing file
file.write_all(content.as_bytes())
.await
.with_context(|| format!("Failed to write prompt: {name}"))?;
// Ensure data is flushed to disk
file.flush()
.await
.with_context(|| format!("Failed to flush prompt: {name}"))?;
// Sync to disk for durability (survive power loss)
file.sync_all()
.await
.with_context(|| format!("Failed to sync prompt to disk: {name}"))?;
// Invalidate cache after successful write
self.invalidate_cache(name).await;
Ok(())
}
Err(e) if e.kind() == ErrorKind::NotFound => {
// File doesn't exist - provide helpful error message
anyhow::bail!("Prompt '{name}' not found. Use add_prompt to create.")
}
Err(e) => {
// Other IO error (permissions, disk full, etc.)
Err(e).with_context(|| format!("Failed to update prompt: {name}"))?
}
}
}
/// Delete a prompt (async)
pub async fn delete_prompt(&self, name: &str) -> Result<()> {
validate_prompt_name(name)?;
let path = self.prompts_dir.join(format!("{name}.j2.md"));
// Attempt delete directly, handle errors appropriately
match fs::remove_file(&path).await {
Ok(()) => {
self.invalidate_cache(name).await;
Ok(())
}
Err(e) if e.kind() == ErrorKind::NotFound => {
anyhow::bail!("Prompt '{name}' not found")
}
Err(e) if e.kind() == ErrorKind::IsADirectory => {
anyhow::bail!("'{name}' is a directory, not a prompt file")
}
Err(e) if e.kind() == ErrorKind::PermissionDenied => {
anyhow::bail!("Permission denied to delete prompt '{name}'")
}
Err(e) => Err(e).with_context(|| format!("Failed to delete prompt: {name}"))?
}
}
/// Render a prompt with parameters (async)
pub async fn render_prompt(
&self,
name: &str,
parameters: Option<HashMap<String, serde_json::Value>>,
) -> Result<String> {
let template = self.load_prompt(name).await?;
render_template(&template, parameters.as_ref()).await
}
/// Invalidate cached entry for a specific prompt
async fn invalidate_cache(&self, name: &str) {
let mut cache = self.cache.write().await;
cache.remove(name);
}
}
/// Get the prompts directory path
/// Supports both local (.kodegen/prompts/) and user-global config with precedence
fn get_prompts_directory() -> Result<PathBuf> {
KodegenConfig::local_config_dir()
.ok()
.map(|dir| dir.join("prompts"))
.or_else(|| {
KodegenConfig::user_config_dir()
.ok()
.map(|dir| dir.join("prompts"))
})
.ok_or_else(|| anyhow::anyhow!("Cannot determine prompts directory"))
}
/// Validate prompt name to prevent path traversal
fn validate_prompt_name(name: &str) -> Result<()> {
// Only alphanumeric, hyphen, underscore
if !name
.chars()
.all(|c| c.is_alphanumeric() || c == '-' || c == '_')
{
anyhow::bail!(
"Invalid prompt name: '{name}'. Only alphanumeric characters, hyphens, and underscores allowed."
);
}
// No path traversal
if name.contains('/') || name.contains('\\') || name.contains("..") {
anyhow::bail!("Invalid prompt name: '{name}'. Path separators and '..' not allowed.");
}
Ok(())
}
/// Quick validation check for prompt names (inline version for list_prompts)
/// Mirrors the logic in validate_prompt_name() for early filtering
fn is_valid_prompt_name(name: &str) -> bool {
!name.is_empty()
&& name
.chars()
.all(|c| c.is_alphanumeric() || c == '-' || c == '_')
&& !name.contains("..")
}
/// Initialize default prompts on first run (async)
async fn initialize_default_prompts(prompts_dir: &Path) -> Result<()> {
// Fast check: does the first default prompt exist?
// If it exists, assume initialization already happened
let first_default = defaults::DEFAULT_PROMPTS[0].0;
let path = prompts_dir.join(format!("{first_default}.j2.md"));
// Check existence, propagating errors instead of masking them
match fs::try_exists(&path).await {
Ok(true) => {
// Initialization already complete - return immediately
return Ok(());
}
Ok(false) => {
// First run: write all default prompts
defaults::write_default_prompts(prompts_dir).await?;
info!(
"Initialized {} default prompts in {}",
defaults::DEFAULT_PROMPTS.len(),
prompts_dir.display()
);
}
Err(e) => {
// Permission or IO error checking existence - propagate with context
return Err(e).with_context(|| {
format!(
"Failed to check if default prompts exist in {}. \
Check directory permissions.",
prompts_dir.display()
)
});
}
}
Ok(())
}