cobolx 1.0.3 → 1.0.4

package/src/agent/fs_agent.rs

@@ -1,11 +1,28 @@
 use super::AgentRouter;
+use super::skills::{AgentKind, append_agent_skills};
 use super::types::merge_tool_call_deltas;
 use super::types::{
     ChatMessage, ChatRequest, FunctionDefinition, StreamOptions, Tool, ToolCall, Usage,
+    WriteBuffer, WriteBufferEntry,
 };
 use crate::memory::MemoryStore;
+use crate::path_safety::{
+    validate_and_resolve_write, validate_sandbox_path as resolve_sandbox_path, write_validated_path,
+};
 use std::path::Path;
 
+fn truncate_utf8_preview(content: &str, max_bytes: usize) -> &str {
+    if content.len() <= max_bytes {
+        return content;
+    }
+
+    let mut end = max_bytes;
+    while end > 0 && !content.is_char_boundary(end) {
+        end -= 1;
+    }
+    &content[..end]
+}
+
 impl AgentRouter {
     /// Validates `user_path` resolves inside `sandbox`.
     /// Returns the canonical absolute path or an error string.
@@ -13,70 +30,7 @@ impl AgentRouter {
         sandbox: &Path,
         user_path: &str,
     ) -> Result<std::path::PathBuf, String> {
-        let normalized = if std::path::Path::new(user_path).is_absolute() {
-            user_path.to_string()
-        } else {
-            user_path.trim_start_matches(['/', '\\']).to_string()
-        };
-
-        let candidate = if std::path::Path::new(&normalized).is_absolute() {
-            std::path::PathBuf::from(&normalized)
-        } else {
-            sandbox.join(&normalized)
-        };
-
-        let sandbox_canon = sandbox
-            .canonicalize()
-            .map_err(|e| format!("Sandbox path error: {e}"))?;
-
-        let clean_canon = |p: &Path| -> String {
-            let s = p.to_string_lossy().into_owned();
-            let s_stripped = if let Some(stripped) = s.strip_prefix(r"\\?\") {
-                stripped.to_string()
-            } else {
-                s
-            };
-            s_stripped.replace('\\', "/").to_lowercase()
-        };
-
-        let sandbox_canon_str = clean_canon(&sandbox_canon);
-
-        let mut existing = candidate.clone();
-        let mut suffix = std::path::PathBuf::new();
-        loop {
-            if existing.exists() {
-                break;
-            }
-            if let Some(parent) = existing.parent() {
-                if let Some(file_name) = existing.file_name() {
-                    suffix = std::path::Path::new(file_name).join(&suffix);
-                    existing = parent.to_path_buf();
-                } else {
-                    break;
-                }
-            } else {
-                break;
-            }
-        }
-
-        let canon_existing = existing
-            .canonicalize()
-            .map_err(|e| format!("Path resolution error: {e}"))?;
-        let resolved = canon_existing.join(&suffix);
-        let resolved_str = clean_canon(&resolved);
-
-        let is_sub = resolved_str == sandbox_canon_str
-            || (resolved_str.starts_with(&sandbox_canon_str)
-                && (sandbox_canon_str.ends_with('/')
-                    || resolved_str.chars().nth(sandbox_canon_str.chars().count()) == Some('/')));
-
-        if !is_sub {
-            return Err(format!(
-                "Access denied: '{}' is outside the sandbox directory",
-                user_path
-            ));
-        }
-        Ok(resolved)
+        resolve_sandbox_path(sandbox, user_path)
     }
 
     /// Phase 1 — silent read-only data retrieval (DB + files).
@@ -110,7 +64,7 @@ impl AgentRouter {
 
         if let Some(first_msg) = messages.get_mut(0) {
             if first_msg.role == "system" {
-                first_msg.content = Some(format!(
+                let mut system_prompt = format!(
                     "You are the COBOLX Filesystem Retrieval Agent. Your ONLY job is to collect \
                     raw data about COBOL files using the tools below. Do NOT explain or interpret \
                     — just gather and output a structured data summary.\n\
@@ -120,7 +74,9 @@ impl AgentRouter {
                     \n\
                     WORKFLOW:\n\
                     1. query_sqlite: SELECT id, path, kind FROM files\n\
-                    2. query_sqlite: get programs, data_items, call_edges, copybook_uses\n\
+                    2. query_sqlite: get programs, data_items, call_edges, copybook_uses, \
+                       program_features, code_blocks, external_ops, identifiers, literals, \
+                       copybook_features\n\
                     3. read_file: raw source text only when needed\n\
                     4. list_directory / search_in_file: locate files if needed\n\
                     \n\
@@ -133,8 +89,20 @@ impl AgentRouter {
                     2. programs(id, name, file_id)\n\
                     3. copybook_uses(id, from_file_id, copybook_name, resolve_status)\n\
                     4. call_edges(id, caller_program_id, callee_name, kind)\n\
-                    5. data_items(id, program_id, name, level, parent_name, pic, usage_clause, section)"
-                ));
+                    5. data_items(id, program_id, name, level, parent_name, pic, usage_clause, section)\n\
+                    6. program_features(program_id, incoming_call_count, outgoing_call_count, is_entrypoint, paragraph_count, external_op_count)\n\
+                    7. code_blocks(program_id, name, kind, parent_section, sequence_no, statement_count)\n\
+                    8. external_ops(program_id, kind, verb, target)\n\
+                    9. identifiers(program_id, kind, value, occurrences)\n\
+                    10. literals(program_id, kind, value, occurrences)\n\
+                    11. copybook_features(copybook_file_id, copybook_name, used_by_program_count, contains_header_fields, contains_error_fields)"
+                );
+                append_agent_skills(
+                    &mut system_prompt,
+                    sandbox_path,
+                    AgentKind::FilesystemRetrieval,
+                )?;
+                first_msg.content = Some(system_prompt);
             }
         }
 
@@ -253,12 +221,20 @@ impl AgentRouter {
                 function: FunctionDefinition {
                     name: "query_sqlite".to_string(),
                     description:
-                        "Run a read-only SELECT query against the project SQLite database \
-                        (files, programs, data_items, call_edges, copybook_uses)."
+                        "Run one read-only SELECT query against the indexed project SQLite database. \
+                        Use this for project facts from files, programs, data_items, call_edges, \
+                        copybook_uses, program_features, code_blocks, external_ops, identifiers, \
+                        literals, or copybook_features. Do not use it for writes, DDL, or guessed \
+                        values."
                             .to_string(),
                     parameters: serde_json::json!({
                         "type": "object",
-                        "properties": { "sql": { "type": "string" } },
+                        "properties": {
+                            "sql": {
+                                "type": "string",
+                                "description": "A single SQLite SELECT statement that reads indexed project data."
+                            }
+                        },
                         "required": ["sql"]
                     }),
                 },
@@ -267,10 +243,17 @@ impl AgentRouter {
                 r#type: "function".to_string(),
                 function: FunctionDefinition {
                     name: "read_file".to_string(),
-                    description: "Read the full text of a sandbox file.".to_string(),
+                    description:
+                        "Read the full text of one sandbox file. Use this when exact source content matters, and pass only a path inside the sandbox."
+                            .to_string(),
                     parameters: serde_json::json!({
                         "type": "object",
-                        "properties": { "path": { "type": "string" } },
+                        "properties": {
+                            "path": {
+                                "type": "string",
+                                "description": "Relative path to one file inside the sandbox."
+                            }
+                        },
                         "required": ["path"]
                     }),
                 },
@@ -279,12 +262,20 @@ impl AgentRouter {
                 r#type: "function".to_string(),
                 function: FunctionDefinition {
                     name: "list_directory".to_string(),
-                    description: "List entries in a sandbox directory.".to_string(),
+                    description:
+                        "List entries in one sandbox directory, optionally filtered by extension. Use this to discover candidate files before reading them."
+                            .to_string(),
                     parameters: serde_json::json!({
                         "type": "object",
                         "properties": {
-                            "path": { "type": "string" },
-                            "extension": { "type": "string" }
+                            "path": {
+                                "type": "string",
+                                "description": "Relative path to a directory inside the sandbox."
+                            },
+                            "extension": {
+                                "type": "string",
+                                "description": "Optional extension filter such as .cbl or .cpy."
+                            }
                         },
                         "required": ["path"]
                     }),
@@ -294,13 +285,20 @@ impl AgentRouter {
                 r#type: "function".to_string(),
                 function: FunctionDefinition {
                     name: "search_in_file".to_string(),
-                    description: "Search for a text pattern (case-insensitive) in a file."
-                        .to_string(),
+                    description:
+                        "Search one sandbox file for a plain-text pattern, case-insensitive, and return matching lines with line numbers."
+                            .to_string(),
                     parameters: serde_json::json!({
                         "type": "object",
                         "properties": {
-                            "path": { "type": "string" },
-                            "pattern": { "type": "string" }
+                            "path": {
+                                "type": "string",
+                                "description": "Relative path to one file inside the sandbox."
+                            },
+                            "pattern": {
+                                "type": "string",
+                                "description": "Plain-text pattern to search for."
+                            }
                         },
                         "required": ["path", "pattern"]
                     }),
@@ -322,8 +320,15 @@ impl AgentRouter {
                 let _ = tx.send("\x01STATUS:Querying project database...".to_string());
                 match MemoryStore::open_or_create(sandbox_path) {
                     Err(e) => serde_json::json!({ "error": format!("DB error: {e}") }).to_string(),
-                    Ok(store) => match store.query_readonly(sql) {
-                        Ok(val) => val.to_string(),
+                    Ok(store) => match store.project_index_is_empty() {
+                        Ok(true) => serde_json::json!({
+                            "error": "Project index is empty. Run /init before asking for indexed project data."
+                        })
+                        .to_string(),
+                        Ok(false) => match store.query_readonly(sql) {
+                            Ok(val) => val.to_string(),
+                            Err(e) => serde_json::json!({ "error": e.to_string() }).to_string(),
+                        },
                         Err(e) => serde_json::json!({ "error": e.to_string() }).to_string(),
                     },
                 }
@@ -338,10 +343,11 @@ impl AgentRouter {
                         Ok(content) => {
                             const MAX: usize = 120_000;
                             let body = if content.len() > MAX {
+                                let preview = truncate_utf8_preview(&content, MAX);
                                 format!(
                                     "[truncated: first {MAX} of {} bytes]\n{}",
                                     content.len(),
-                                    &content[..MAX]
+                                    preview
                                 )
                             } else {
                                 content
@@ -423,4 +429,126 @@ impl AgentRouter {
             }
         })
     }
+
+    /// Writes a file to the sandbox. If a buffer is provided, it is pushed to the buffer instead of writing physically.
+    /// Returns the resolved path or an error string.
+    pub(crate) fn write_file(
+        &self,
+        sandbox: &Path,
+        user_path: &str,
+        content: &str,
+        buffer: Option<&WriteBuffer>,
+    ) -> Result<std::path::PathBuf, String> {
+        let full_path = validate_and_resolve_write(sandbox, user_path)?;
+        if let Some(buf) = buffer {
+            if let Ok(mut lock) = buf.lock() {
+                lock.push((full_path.clone(), content.to_string()));
+            } else {
+                return Err("Failed to lock write buffer".to_string());
+            }
+        } else {
+            write_validated_path(&full_path, content)?;
+        }
+        Ok(full_path)
+    }
+
+    /// Commits a list of buffered writes to disk.
+    pub(crate) fn commit_write_buffer(&self, buffer: &[WriteBufferEntry]) -> Result<(), String> {
+        for (full_path, content) in buffer {
+            write_validated_path(full_path, content)?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agent::types::{FunctionCall, ToolCall};
+    use std::io::Write;
+
+    #[test]
+    fn readonly_tools_descriptions_include_usage_constraints() {
+        let tools = AgentRouter::build_readonly_tools();
+
+        let query_sqlite = tools
+            .iter()
+            .find(|t| t.function.name == "query_sqlite")
+            .unwrap();
+        assert!(query_sqlite.function.description.contains("SELECT"));
+        assert!(query_sqlite.function.description.contains("read-only"));
+
+        let read_file = tools
+            .iter()
+            .find(|t| t.function.name == "read_file")
+            .unwrap();
+        assert!(read_file.function.description.contains("sandbox"));
+        assert!(read_file.function.description.contains("full text"));
+
+        let list_directory = tools
+            .iter()
+            .find(|t| t.function.name == "list_directory")
+            .unwrap();
+        assert!(list_directory.function.description.contains("directory"));
+        assert!(list_directory.function.description.contains("extension"));
+
+        let search_in_file = tools
+            .iter()
+            .find(|t| t.function.name == "search_in_file")
+            .unwrap();
+        assert!(
+            search_in_file
+                .function
+                .description
+                .contains("case-insensitive")
+        );
+        assert!(search_in_file.function.description.contains("line"));
+    }
+
+    #[tokio::test]
+    async fn read_file_truncation_handles_utf8_boundaries_without_panicking() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("utf8.cbl");
+        let mut file = std::fs::File::create(&path).unwrap();
+        let content = format!("a{}", "你".repeat(40_100));
+        file.write_all(content.as_bytes()).unwrap();
+
+        let tc = ToolCall {
+            id: "1".to_string(),
+            r#type: "function".to_string(),
+            function: FunctionCall {
+                name: "read_file".to_string(),
+                arguments: serde_json::json!({ "path": "utf8.cbl" }).to_string(),
+            },
+        };
+        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
+
+        let result = AgentRouter::execute_readonly_tool(&tc, dir.path(), &tx).await;
+        assert!(result.is_ok());
+
+        let result_json = result.unwrap();
+        let parsed: serde_json::Value = serde_json::from_str(&result_json).unwrap();
+        let body = parsed["content"].as_str().unwrap_or("");
+        assert!(body.contains("[truncated:"), "tool result: {}", result_json);
+    }
+
+    #[tokio::test]
+    async fn query_sqlite_returns_init_guidance_when_index_is_empty() {
+        let dir = tempfile::tempdir().unwrap();
+        let tc = ToolCall {
+            id: "2".to_string(),
+            r#type: "function".to_string(),
+            function: FunctionCall {
+                name: "query_sqlite".to_string(),
+                arguments: serde_json::json!({ "sql": "SELECT * FROM files" }).to_string(),
+            },
+        };
+        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
+
+        let result = AgentRouter::execute_readonly_tool(&tc, dir.path(), &tx)
+            .await
+            .unwrap();
+        assert!(result.contains("/init"));
+        assert!(result.contains("index"));
+    }
 }