Merge PR #236: fix(cortex-cli): warn about missing agent references during session import

Author: factorydroid

Cargo.lock

@@ -68,6 +68,9 @@ tar = "0.4"
 # For scrape command (HTML parsing)
 scraper = "0.22"
 
+# For agent reference extraction from messages
+regex = { workspace = true }
+
 # For mDNS service discovery
 hostname = { workspace = true }
 

cortex-cli/Cargo.toml

@@ -323,7 +323,7 @@ fn get_project_agents_dirs() -> Vec<PathBuf> {
 }
 
 /// Load all agents from various sources.
-fn load_all_agents() -> Result<Vec<AgentInfo>> {
+pub fn load_all_agents() -> Result<Vec<AgentInfo>> {
     let mut agents = Vec::new();
     let mut seen_names = std::collections::HashSet::new();
 

cortex-cli/src/agent_cmd.rs

@@ -55,6 +55,12 @@ pub struct SessionMetadata {
     /// Model used for the session.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub model: Option<String>,
+    /// Agent used for the session (if any).
+    #[serde(skip_serializing_if = "Option::is_none", default)]
+    pub agent: Option<String>,
+    /// Agent references found in messages (for validation during import).
+    #[serde(skip_serializing_if = "Option::is_none", default)]
+    pub agent_refs: Option<Vec<String>>,
 }
 
 /// Message in export format.
@@ -147,6 +153,13 @@ impl ExportCommand {
 
         // Extract metadata
         let meta = get_session_meta(&entries);
+
+        // Extract messages from events
+        let messages = extract_messages(&entries);
+
+        // Extract agent references from messages (@agent mentions)
+        let agent_refs = extract_agent_refs(&messages);
+
         let session_meta = SessionMetadata {
             id: conversation_id.to_string(),
             title: derive_title(&entries),
@@ -155,11 +168,14 @@ impl ExportCommand {
                 .unwrap_or_else(|| chrono::Utc::now().to_rfc3339()),
             cwd: meta.map(|m| m.cwd.clone()),
             model: meta.and_then(|m| m.model.clone()),
+            agent: None, // TODO: Extract from session config if stored
+            agent_refs: if agent_refs.is_empty() {
+                None
+            } else {
+                Some(agent_refs)
+            },
         };
 
-        // Extract messages from events
-        let messages = extract_messages(&entries);
-
         // Build export
         let export = SessionExport {
             version: 1,
@@ -292,6 +308,29 @@ fn extract_messages(
     messages
 }
 
+/// Extract agent references (@mentions) from messages.
+/// Returns a deduplicated list of agent names that are referenced.
+fn extract_agent_refs(messages: &[ExportMessage]) -> Vec<String> {
+    use std::collections::HashSet;
+
+    // Regex to match @agent mentions (e.g., @explore, @general, @my-custom-agent)
+    let re = regex::Regex::new(r"@([a-zA-Z][a-zA-Z0-9_-]*)").unwrap();
+
+    let mut agent_refs: HashSet<String> = HashSet::new();
+
+    for message in messages {
+        for cap in re.captures_iter(&message.content) {
+            if let Some(agent_name) = cap.get(1) {
+                agent_refs.insert(agent_name.as_str().to_string());
+            }
+        }
+    }
+
+    let mut refs: Vec<String> = agent_refs.into_iter().collect();
+    refs.sort();
+    refs
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -306,6 +345,8 @@ mod tests {
                 created_at: "2024-01-01T00:00:00Z".to_string(),
                 cwd: Some("/home/user".to_string()),
                 model: Some("claude-3".to_string()),
+                agent: None,
+                agent_refs: None,
             },
             messages: vec![
                 ExportMessage {
@@ -330,4 +371,54 @@ mod tests {
         assert!(json.contains("\"role\": \"user\""));
         assert!(json.contains("\"content\": \"Hello\""));
     }
+
+    #[test]
+    fn test_extract_agent_refs() {
+        let messages = vec![
+            ExportMessage {
+                role: "user".to_string(),
+                content: "@explore find the main function".to_string(),
+                tool_calls: None,
+                tool_call_id: None,
+                timestamp: None,
+            },
+            ExportMessage {
+                role: "user".to_string(),
+                content: "@research analyze this code @my-agent".to_string(),
+                tool_calls: None,
+                tool_call_id: None,
+                timestamp: None,
+            },
+        ];
+
+        let refs = extract_agent_refs(&messages);
+        assert_eq!(refs.len(), 3);
+        assert!(refs.contains(&"explore".to_string()));
+        assert!(refs.contains(&"research".to_string()));
+        assert!(refs.contains(&"my-agent".to_string()));
+    }
+
+    #[test]
+    fn test_extract_agent_refs_no_duplicates() {
+        let messages = vec![
+            ExportMessage {
+                role: "user".to_string(),
+                content: "@explore task 1".to_string(),
+                tool_calls: None,
+                tool_call_id: None,
+                timestamp: None,
+            },
+            ExportMessage {
+                role: "user".to_string(),
+                content: "@explore task 2".to_string(),
+                tool_calls: None,
+                tool_call_id: None,
+                timestamp: None,
+            },
+        ];
+
+        let refs = extract_agent_refs(&messages);
+        assert_eq!(refs.len(), 1);
+        assert!(refs.contains(&"explore".to_string()));
+    }
 }

cortex-cli/src/export_cmd.rs

@@ -4,6 +4,7 @@
 
 use anyhow::{Context, Result, bail};
 use clap::Parser;
+use std::collections::HashSet;
 use std::path::PathBuf;
 
 use crate::styled_output::{print_info, print_success, print_warning};
@@ -14,6 +15,7 @@ use cortex_protocol::{
     ParsedCommand, UserMessageEvent,
 };
 
+use crate::agent_cmd::load_all_agents;
 use crate::export_cmd::{ExportMessage, SessionExport};
 
 /// Import a session from JSON format.
@@ -102,6 +104,25 @@ impl ImportCommand {
             );
         }
 
+        // Validate agent references in the imported session
+        let missing_agents = validate_agent_references(&export)?;
+        if !missing_agents.is_empty() {
+            eprintln!(
+                "Warning: The following agent references in this session are not available locally:"
+            );
+            for agent in &missing_agents {
+                eprintln!("  - @{}", agent);
+            }
+            eprintln!();
+            eprintln!(
+                "The session will be imported, but agent-related functionality may not work as expected."
+            );
+            eprintln!(
+                "To fix this, create the missing agents using 'cortex agent create <name>' or copy them from the source system."
+            );
+            eprintln!();
+        }
+
         // Generate a new session ID (we always create a new session on import)
         let new_conversation_id = ConversationId::new();
 
@@ -235,6 +256,54 @@ async fn fetch_url(url: &str) -> Result<String> {
     }
 }
 
+/// Validate agent references in the imported session.
+/// Returns a list of missing agent names that are referenced but not available locally.
+fn validate_agent_references(export: &SessionExport) -> Result<Vec<String>> {
+    // Get all locally available agents
+    let local_agents: HashSet<String> = match load_all_agents() {
+        Ok(agents) => agents.into_iter().map(|a| a.name).collect(),
+        Err(e) => {
+            // If we can't load agents, log a warning but continue
+            tracing::warn!("Could not load local agents for validation: {}", e);
+            HashSet::new()
+        }
+    };
+
+    let mut missing_agents = Vec::new();
+
+    // Check the session's agent field (if the session was started with -a <agent>)
+    if let Some(ref agent_name) = export.session.agent {
+        if !local_agents.contains(agent_name) {
+            missing_agents.push(agent_name.clone());
+        }
+    }
+
+    // Check agent_refs from the export metadata (pre-extracted @mentions)
+    if let Some(ref agent_refs) = export.session.agent_refs {
+        for agent_ref in agent_refs {
+            if !local_agents.contains(agent_ref) && !missing_agents.contains(agent_ref) {
+                missing_agents.push(agent_ref.clone());
+            }
+        }
+    }
+
+    // Also scan messages for @agent mentions (in case they weren't pre-extracted)
+    let re = regex::Regex::new(r"@([a-zA-Z][a-zA-Z0-9_-]*)").unwrap();
+    for message in &export.messages {
+        for cap in re.captures_iter(&message.content) {
+            if let Some(agent_name) = cap.get(1) {
+                let name = agent_name.as_str().to_string();
+                if !local_agents.contains(&name) && !missing_agents.contains(&name) {
+                    missing_agents.push(name);
+                }
+            }
+        }
+    }
+
+    missing_agents.sort();
+    Ok(missing_agents)
+}
+
 /// Convert an export message to a protocol event.
 fn message_to_event(message: &ExportMessage, turn_id: &mut u64, cwd: &PathBuf) -> Result<Event> {
     let event_msg = match message.role.as_str() {
@@ -422,4 +491,79 @@ mod tests {
         assert_eq!(preview_len, 200);
         assert_eq!(&long_content[..preview_len].len(), &200);
     }
+
+    #[test]
+    fn test_validate_agent_references_with_missing_agents() {
+        use crate::export_cmd::SessionMetadata;
+
+        // Create an export with agent references that don't exist locally
+        let export = SessionExport {
+            version: 1,
+            session: SessionMetadata {
+                id: "test-123".to_string(),
+                title: None,
+                created_at: "2024-01-01T00:00:00Z".to_string(),
+                cwd: None,
+                model: None,
+                agent: Some("custom-nonexistent-agent".to_string()),
+                agent_refs: Some(vec!["another-nonexistent".to_string()]),
+            },
+            messages: vec![ExportMessage {
+                role: "user".to_string(),
+                content: "@yet-another-missing help me".to_string(),
+                tool_calls: None,
+                tool_call_id: None,
+                timestamp: None,
+            }],
+        };
+
+        let missing = validate_agent_references(&export).unwrap();
+
+        // Should find at least the explicitly nonexistent ones
+        // (built-in agents like 'explore', 'research' will exist)
+        assert!(missing.contains(&"custom-nonexistent-agent".to_string()));
+        assert!(missing.contains(&"another-nonexistent".to_string()));
+        assert!(missing.contains(&"yet-another-missing".to_string()));
+    }
+
+    #[test]
+    fn test_validate_agent_references_with_builtin_agents() {
+        use crate::export_cmd::SessionMetadata;
+
+        // Create an export referencing only built-in agents
+        let export = SessionExport {
+            version: 1,
+            session: SessionMetadata {
+                id: "test-123".to_string(),
+                title: None,
+                created_at: "2024-01-01T00:00:00Z".to_string(),
+                cwd: None,
+                model: None,
+                agent: None,
+                agent_refs: None,
+            },
+            messages: vec![
+                ExportMessage {
+                    role: "user".to_string(),
+                    content: "@build help me compile".to_string(),
+                    tool_calls: None,
+                    tool_call_id: None,
+                    timestamp: None,
+                },
+                ExportMessage {
+                    role: "user".to_string(),
+                    content: "@plan create a plan".to_string(),
+                    tool_calls: None,
+                    tool_call_id: None,
+                    timestamp: None,
+                },
+            ],
+        };
+
+        let missing = validate_agent_references(&export).unwrap();
+
+        // Built-in agents should not be reported as missing
+        assert!(!missing.contains(&"build".to_string()));
+        assert!(!missing.contains(&"plan".to_string()));
+    }
 }