Merge PR #252: fix: batch fixes for issues #2363-#2379

Author: factorydroid

cortex-cli/src/debug_cmd.rs

@@ -149,13 +149,9 @@ async fn run_config(args: ConfigArgs) -> Result<()> {
         ];
         for var in cortex_vars {
             if let Ok(val) = std::env::var(var) {
-                // Mask API keys
-                let display_val = if var.contains("API_KEY") || var.contains("SECRET") {
-                    if val.len() > 8 {
-                        format!("{}...{}", &val[..4], &val[val.len() - 4..])
-                    } else {
-                        "***".to_string()
-                    }
+                // Mask sensitive values (API keys, secrets, tokens, passwords, credentials)
+                let display_val = if is_sensitive_var_name(var) {
+                    redact_sensitive_value(&val)
                 } else {
                     val
                 };
@@ -2067,6 +2063,40 @@ impl DebugCli {
     }
 }
 
+/// Patterns that indicate a variable contains sensitive data.
+const SENSITIVE_PATTERNS: &[&str] = &[
+    "API_KEY",
+    "SECRET",
+    "TOKEN",
+    "PASSWORD",
+    "CREDENTIAL",
+    "PRIVATE",
+    "AUTH",
+    "ACCESS_KEY",
+    "BEARER",
+    "SESSION",
+];
+
+/// Check if an environment variable name indicates sensitive data.
+fn is_sensitive_var_name(name: &str) -> bool {
+    let name_upper = name.to_uppercase();
+    SENSITIVE_PATTERNS
+        .iter()
+        .any(|pattern| name_upper.contains(pattern))
+}
+
+/// Redact a sensitive value, showing only first and last few characters.
+fn redact_sensitive_value(value: &str) -> String {
+    if value.is_empty() {
+        return "[EMPTY]".to_string();
+    }
+    if value.len() <= 8 {
+        return "[REDACTED]".to_string();
+    }
+    // Show first 4 and last 4 characters
+    format!("{}...{}", &value[..4], &value[value.len() - 4..])
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -2086,4 +2116,38 @@ mod tests {
         assert_eq!(format_size(1048576), "1.00 MB");
         assert_eq!(format_size(1073741824), "1.00 GB");
     }
+
+    #[test]
+    fn test_is_sensitive_var_name() {
+        // Should match sensitive patterns
+        assert!(is_sensitive_var_name("OPENAI_API_KEY"));
+        assert!(is_sensitive_var_name("DATABASE_PASSWORD"));
+        assert!(is_sensitive_var_name("AWS_SECRET_ACCESS_KEY"));
+        assert!(is_sensitive_var_name("AUTH_TOKEN"));
+        assert!(is_sensitive_var_name("GITHUB_TOKEN"));
+        assert!(is_sensitive_var_name("PRIVATE_KEY"));
+        assert!(is_sensitive_var_name("CREDENTIAL_FILE"));
+        assert!(is_sensitive_var_name("BEARER_TOKEN"));
+
+        // Should not match non-sensitive patterns
+        assert!(!is_sensitive_var_name("PATH"));
+        assert!(!is_sensitive_var_name("HOME"));
+        assert!(!is_sensitive_var_name("USER"));
+        assert!(!is_sensitive_var_name("EDITOR"));
+        assert!(!is_sensitive_var_name("SHELL"));
+    }
+
+    #[test]
+    fn test_redact_sensitive_value() {
+        // Empty value
+        assert_eq!(redact_sensitive_value(""), "[EMPTY]");
+
+        // Short value (8 or fewer chars)
+        assert_eq!(redact_sensitive_value("short"), "[REDACTED]");
+        assert_eq!(redact_sensitive_value("12345678"), "[REDACTED]");
+
+        // Longer value shows first/last 4 chars
+        assert_eq!(redact_sensitive_value("sk-abc123xyz789"), "sk-a...9789");
+        assert_eq!(redact_sensitive_value("supersecretpassword"), "supe...word");
+    }
 }

cortex-cli/src/run_cmd.rs

@@ -152,6 +152,11 @@ pub struct RunCli {
     /// Timeout in seconds (0 for no timeout).
     #[arg(long = "timeout", default_value_t = 0)]
     pub timeout: u64,
+
+    /// Preview what would be sent without executing.
+    /// Shows estimated token counts including system prompt and tool definitions.
+    #[arg(long = "dry-run")]
+    pub dry_run: bool,
 }
 
 /// Tool display information for formatted output.
@@ -467,6 +472,11 @@ impl RunCli {
         attachments: &[FileAttachment],
         session_mode: SessionMode,
     ) -> Result<()> {
+        // Handle dry-run mode - show token estimates without executing
+        if self.dry_run {
+            return self.run_dry_run(message, attachments).await;
+        }
+
         let is_json = matches!(self.format, OutputFormat::Json | OutputFormat::Jsonl);
         let is_terminal = io::stdout().is_terminal();
 
@@ -828,6 +838,109 @@ impl RunCli {
 
         Ok(())
     }
+
+    /// Run in dry-run mode - show token estimates without executing.
+    async fn run_dry_run(&self, message: &str, attachments: &[FileAttachment]) -> Result<()> {
+        use cortex_engine::tokenizer::{TokenCounter, TokenizerType};
+
+        let config = cortex_engine::Config::default();
+        let model = self
+            .model
+            .as_ref()
+            .map(|m| resolve_model_alias(m).to_string())
+            .unwrap_or_else(|| config.model.clone());
+
+        let mut counter = TokenCounter::for_model(&model);
+
+        // Count user prompt tokens
+        let user_prompt_tokens = counter.count(message);
+
+        // Count attachment tokens
+        let mut attachment_tokens = 0u32;
+        for attachment in attachments {
+            let content =
+                std::fs::read_to_string(&attachment.path).unwrap_or_else(|_| String::new());
+            attachment_tokens += counter.count(&content);
+            // Add overhead for file markers
+            attachment_tokens += 20; // Approximate overhead for "--- File: ... ---" markers
+        }
+
+        // Estimate system prompt tokens (typical system prompt is ~500-2000 tokens)
+        // This is an approximation as the actual system prompt varies
+        let system_prompt_tokens = 1500u32;
+
+        // Estimate tool definition tokens
+        // Each tool definition is approximately 100-200 tokens on average
+        // Common tools: Execute, Read, Write, Edit, LS, Grep, Glob, etc.
+        let tool_count = 15; // Approximate number of default tools
+        let tool_tokens = tool_count * 150; // ~150 tokens per tool definition
+
+        // Calculate totals
+        let total_input_tokens =
+            user_prompt_tokens + attachment_tokens + system_prompt_tokens + tool_tokens;
+
+        // Output based on format
+        if matches!(self.format, OutputFormat::Json | OutputFormat::Jsonl) {
+            let output = serde_json::json!({
+                "dry_run": true,
+                "model": model,
+                "token_estimates": {
+                    "user_prompt": user_prompt_tokens,
+                    "attachments": attachment_tokens,
+                    "system_prompt": system_prompt_tokens,
+                    "tool_definitions": tool_tokens,
+                    "total_input": total_input_tokens,
+                },
+                "message_preview": if message.len() > 100 {
+                    format!("{}...", &message[..100])
+                } else {
+                    message.to_string()
+                },
+                "attachment_count": attachments.len(),
+            });
+            println!("{}", serde_json::to_string_pretty(&output)?);
+        } else {
+            println!("Dry Run - Token Estimate");
+            println!("{}", "=".repeat(50));
+            println!();
+            println!("Model: {}", model);
+            println!();
+            println!("Token Breakdown:");
+            println!("  User prompt:      {:>8} tokens", user_prompt_tokens);
+            if !attachments.is_empty() {
+                println!(
+                    "  Attachments:      {:>8} tokens ({} files)",
+                    attachment_tokens,
+                    attachments.len()
+                );
+            }
+            println!(
+                "  System prompt:    {:>8} tokens (estimated)",
+                system_prompt_tokens
+            );
+            println!(
+                "  Tool definitions: {:>8} tokens (estimated, {} tools)",
+                tool_tokens, tool_count
+            );
+            println!("  {}", "-".repeat(30));
+            println!("  Total input:      {:>8} tokens", total_input_tokens);
+            println!();
+            println!("Note: System prompt and tool definition token counts are estimates.");
+            println!("Actual counts may vary based on agent configuration.");
+            if !message.is_empty() {
+                println!();
+                println!("Message preview:");
+                let preview = if message.len() > 200 {
+                    format!("  {}...", &message[..200])
+                } else {
+                    format!("  {}", message)
+                };
+                println!("{}", preview);
+            }
+        }
+
+        Ok(())
+    }
 }
 
 /// Session handling mode.

cortex-cli/src/scrape_cmd.rs

@@ -124,6 +124,13 @@ pub struct ScrapeCommand {
     #[arg(long, default_value = "highest", value_name = "RESOLUTION")]
     pub image_resolution: String,
 
+    /// Attempt to traverse shadow DOM elements.
+    /// When enabled, will look for <template> tags with shadowrootmode attribute
+    /// and include their content. This is a best-effort feature as shadow DOM
+    /// content is typically only accessible via JavaScript execution.
+    #[arg(long)]
+    pub include_shadow_dom: bool,
+
     /// Show verbose output (includes fetching info).
     #[arg(short, long)]
     pub verbose: bool,
@@ -224,7 +231,14 @@ impl ScrapeCommand {
 
     /// Process HTML content based on options.
     fn process_html(&self, html: &str, format: OutputFormat) -> Result<String> {
-        let document = Html::parse_document(html);
+        // Preprocess HTML to extract shadow DOM content if requested
+        let processed_html = if self.include_shadow_dom {
+            extract_shadow_dom_content(html)
+        } else {
+            html.to_string()
+        };
+
+        let document = Html::parse_document(&processed_html);
 
         // If a selector is provided, extract only that content
         let content_html = if let Some(selector_str) = &self.selector {
@@ -269,6 +283,36 @@ impl ScrapeCommand {
     }
 }
 
+/// Extract shadow DOM content from HTML.
+/// This is a best-effort approach that handles declarative shadow DOM (template tags
+/// with shadowrootmode attribute) and replaces custom elements with their shadow content.
+fn extract_shadow_dom_content(html: &str) -> String {
+    let document = Html::parse_document(html);
+    let mut result = html.to_string();
+
+    // Look for declarative shadow DOM templates
+    // These are <template shadowrootmode="open"> or <template shadowroot="open"> tags
+    if let Ok(template_selector) = Selector::parse("template") {
+        for template in document.select(&template_selector) {
+            // Check for shadow root attributes
+            let has_shadow_attr = template.value().attr("shadowrootmode").is_some()
+                || template.value().attr("shadowroot").is_some();
+
+            if has_shadow_attr {
+                // Get the inner HTML of the template
+                let inner_html = template.inner_html();
+
+                // Replace the template with its content
+                // This makes shadow DOM content visible to the scraper
+                let template_html = template.html();
+                result = result.replace(&template_html, &inner_html);
+            }
+        }
+    }
+
+    result
+}
+
 /// Parse custom headers from command line arguments.
 fn parse_headers(headers: &[String]) -> Result<HashMap<String, String>> {
     let mut result = HashMap::new();

cortex-cli/src/stats_cmd.rs

@@ -104,11 +104,44 @@ pub struct DateRange {
 }
 
 /// Pricing information per 1M tokens.
+#[derive(Debug, Clone)]
 struct ModelPricing {
     input_per_million: f64,
     output_per_million: f64,
 }
 
+/// Custom pricing configuration loaded from config file or environment.
+/// This allows users to override default pricing when provider prices change.
+fn load_custom_pricing() -> std::collections::HashMap<String, ModelPricing> {
+    let mut custom = std::collections::HashMap::new();
+
+    // Try to load from environment variables in format:
+    // CORTEX_PRICING_<MODEL>=<input_price>,<output_price>
+    // Example: CORTEX_PRICING_GPT4O=2.5,10.0
+    for (key, value) in std::env::vars() {
+        if let Some(model_suffix) = key.strip_prefix("CORTEX_PRICING_") {
+            let model_name = model_suffix.to_lowercase().replace('_', "-");
+            let parts: Vec<&str> = value.split(',').collect();
+            if parts.len() == 2 {
+                if let (Ok(input), Ok(output)) = (
+                    parts[0].trim().parse::<f64>(),
+                    parts[1].trim().parse::<f64>(),
+                ) {
+                    custom.insert(
+                        model_name,
+                        ModelPricing {
+                            input_per_million: input,
+                            output_per_million: output,
+                        },
+                    );
+                }
+            }
+        }
+    }
+
+    custom
+}
+
 impl StatsCli {
     /// Run the stats command.
     pub async fn run(self) -> Result<()> {
@@ -165,8 +198,26 @@ fn get_cortex_home() -> PathBuf {
 }
 
 /// Get pricing for a model.
+/// Checks custom pricing from environment first, then falls back to defaults.
 fn get_model_pricing(model: &str) -> ModelPricing {
-    // Pricing per 1M tokens (approximate as of late 2024)
+    // First check for custom pricing from environment
+    let custom_pricing = load_custom_pricing();
+    let model_lower = model.to_lowercase();
+
+    // Check for exact match in custom pricing
+    if let Some(pricing) = custom_pricing.get(&model_lower) {
+        return pricing.clone();
+    }
+
+    // Check for partial match in custom pricing (e.g., "gpt-4o" matches "gpt-4o-mini")
+    for (key, pricing) in &custom_pricing {
+        if model_lower.contains(key) {
+            return pricing.clone();
+        }
+    }
+
+    // Fall back to default pricing (may be outdated - users can override via CORTEX_PRICING_*)
+    // Pricing per 1M tokens (as of late 2024/early 2025 - may change)
     match model {
         // Anthropic
         m if m.contains("claude-opus-4") || m.contains("opus-4") => ModelPricing {

cortex-engine/src/client/types.rs

@@ -15,6 +15,11 @@ pub struct CompletionRequest {
     /// Temperature for sampling.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub temperature: Option<f32>,
+    /// Random seed for reproducibility.
+    /// When set, the same seed with identical inputs should produce deterministic outputs.
+    /// Note: This is applied to all model calls including tool invocations.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub seed: Option<u64>,
     /// Tools available for the model.
     #[serde(skip_serializing_if = "Vec::is_empty")]
     pub tools: Vec<ToolDefinition>,
@@ -30,6 +35,7 @@ impl Default for CompletionRequest {
             model: String::new(),
             max_tokens: None,
             temperature: None,
+            seed: None,
             tools: vec![],
             stream: true,
         }