Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions src/cortex-cli/src/export_cmd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ use cortex_engine::rollout::get_rollout_path;
use cortex_engine::rollout::reader::{RolloutItem, get_session_meta, read_rollout};
use cortex_protocol::{ConversationId, EventMsg};

pub(crate) const AGENT_MENTION_REGEX: &str = r"@([a-zA-Z][a-zA-Z0-9_-]*(?:\.[a-zA-Z0-9_-]+)*)";

/// Export format for sessions.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)]
pub enum ExportFormat {
Expand Down Expand Up @@ -358,8 +360,8 @@ fn extract_messages(
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();
// Regex to match @agent mentions (e.g., @explore, @my-custom-agent, @acme.deploy)
let re = regex::Regex::new(AGENT_MENTION_REGEX).unwrap();

let mut agent_refs: HashSet<String> = HashSet::new();

Expand Down Expand Up @@ -443,6 +445,28 @@ mod tests {
assert!(refs.contains(&"my-agent".to_string()));
}

#[test]
fn test_extract_agent_refs_with_dotted_names() {
let messages = vec![ExportMessage {
role: "user".to_string(),
content: "Run @acme.deploy, then @foo.bar. Keep @my-agent too.".to_string(),
tool_calls: None,
tool_call_id: None,
timestamp: None,
}];

let refs = extract_agent_refs(&messages);
assert_eq!(
refs,
vec![
"acme.deploy".to_string(),
"foo.bar".to_string(),
"my-agent".to_string()
]
);
assert!(!refs.contains(&"acme".to_string()));
}

#[test]
fn test_extract_agent_refs_no_duplicates() {
let messages = vec![
Expand Down
36 changes: 34 additions & 2 deletions src/cortex-cli/src/import_cmd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use cortex_protocol::{
};

use crate::agent_cmd::load_all_agents;
use crate::export_cmd::{ExportMessage, SessionExport};
use crate::export_cmd::{AGENT_MENTION_REGEX, ExportMessage, SessionExport};

/// Maximum depth for processing messages to prevent stack overflow from deeply nested structures.
const MAX_PROCESSING_DEPTH: usize = 10000;
Expand Down Expand Up @@ -289,7 +289,7 @@ fn validate_agent_references(export: &SessionExport) -> Result<Vec<String>> {
}

// 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();
let re = regex::Regex::new(AGENT_MENTION_REGEX).unwrap();
for message in &export.messages {
for cap in re.captures_iter(&message.content) {
if let Some(agent_name) = cap.get(1) {
Expand Down Expand Up @@ -650,6 +650,38 @@ mod tests {
assert!(missing.contains(&"yet-another-missing".to_string()));
}

#[test]
fn test_validate_agent_references_with_dotted_missing_agent() {
use crate::export_cmd::SessionMetadata;

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: "Please run @acme.deploy for me, then @foo.bar.".to_string(),
tool_calls: None,
tool_call_id: None,
timestamp: None,
}],
};

let missing = validate_agent_references(&export).unwrap();

assert!(missing.contains(&"acme.deploy".to_string()));
assert!(missing.contains(&"foo.bar".to_string()));
assert!(!missing.contains(&"acme".to_string()));
assert!(!missing.contains(&"foo".to_string()));
}

#[test]
fn test_validate_agent_references_with_builtin_agents() {
use crate::export_cmd::SessionMetadata;
Expand Down