use std::fs;
use std::sync::Arc;
use tempfile::tempdir;
use ucm_core::{Block, BlockId, Content, Document};
use ucp_agent::{
AgentCapabilities, AgentError, AgentTraversal, ExpandDirection, ExpandOptions, GlobalLimits,
MockRagProvider, RagProvider, SearchOptions, SessionConfig, SessionLimits, ViewMode,
};
use ucp_codegraph::{build_code_graph, CodeGraphBuildInput, CodeGraphExtractorConfig};
fn create_test_document() -> Document {
let mut doc = Document::create();
let root = doc.root;
let child1 =
Block::new(Content::text("Child 1 - Introduction"), Some("heading1")).with_tag("important");
let child2 =
Block::new(Content::text("Child 2 - Methods"), Some("heading2")).with_tag("methods");
let child3 = Block::new(
Content::text("Child 3 - Simple paragraph"),
Some("paragraph"),
);
let child1_id = doc.add_block(child1, &root).unwrap();
let child2_id = doc.add_block(child2, &root).unwrap();
let _child3_id = doc.add_block(child3, &root).unwrap();
let grandchild1a = Block::new(
Content::text("This is the introduction paragraph with important details."),
Some("paragraph"),
)
.with_tag("important")
.with_tag("introduction");
let grandchild1b = Block::new(
Content::code("rust", "fn main() { println!(\"Hello\"); }"),
None,
);
let grandchild2 = Block::new(
Content::text("Methodology section content here."),
Some("paragraph"),
)
.with_tag("methods");
doc.add_block(grandchild1a, &child1_id).unwrap();
doc.add_block(grandchild1b, &child1_id).unwrap();
doc.add_block(grandchild2, &child2_id).unwrap();
doc
}
fn fake_block_id() -> BlockId {
BlockId::from_bytes([
0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
])
}
fn create_codegraph_document() -> Document {
let dir = tempdir().unwrap();
fs::create_dir_all(dir.path().join("src")).unwrap();
fs::write(
dir.path().join("src/util.rs"),
"pub fn util() -> i32 { 1 }\n",
)
.unwrap();
fs::write(
dir.path().join("src/lib.rs"),
"mod util;\npub fn add(a: i32, b: i32) -> i32 { util::util() + a + b }\n",
)
.unwrap();
let repository_path = dir.path().to_path_buf();
std::mem::forget(dir);
build_code_graph(&CodeGraphBuildInput {
repository_path,
commit_hash: "agent-context-tests".to_string(),
config: CodeGraphExtractorConfig::default(),
})
.unwrap()
.document
}
#[test]
fn test_session_creation_and_close() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal
.create_session(SessionConfig::default())
.expect("Should create session");
assert!(!session_id.0.is_nil(), "Session ID should not be nil");
traversal
.close_session(&session_id)
.expect("Should close session");
let result = traversal.close_session(&session_id);
assert!(matches!(result, Err(AgentError::SessionNotFound(_))));
}
#[test]
fn test_session_with_config() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let config = SessionConfig::new()
.with_name("test-session")
.with_view_mode(ViewMode::Preview { length: 50 });
let session_id = traversal
.create_session(config)
.expect("Should create session with config");
assert!(!session_id.0.is_nil());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_max_sessions_limit() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc).with_global_limits(GlobalLimits {
max_sessions: 2,
..Default::default()
});
let _s1 = traversal.create_session(SessionConfig::default()).unwrap();
let _s2 = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.create_session(SessionConfig::default());
assert!(matches!(
result,
Err(AgentError::MaxSessionsReached { max: 2 })
));
}
#[test]
fn test_session_with_custom_capabilities() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let config = SessionConfig::new().with_capabilities(AgentCapabilities::read_only());
let session_id = traversal
.create_session(config)
.expect("Should create read-only session");
let _ = traversal.navigate_to(&session_id, root_id);
let result = traversal.context_add(&session_id, root_id, None, None);
assert!(matches!(
result,
Err(AgentError::OperationNotPermitted { .. })
));
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_navigate_to_root() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.navigate_to(&session_id, root_id);
assert!(result.is_ok());
let nav_result = result.unwrap();
assert_eq!(nav_result.position, root_id);
assert!(nav_result.refreshed);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_navigate_to_nonexistent_block() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let fake_id = fake_block_id();
let result = traversal.navigate_to(&session_id, fake_id);
assert!(matches!(result, Err(AgentError::BlockNotFound(_))));
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_go_back_empty_history() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.go_back(&session_id, 1);
assert!(matches!(result, Err(AgentError::EmptyHistory)));
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_navigate_and_go_back() {
let doc = create_test_document();
let root_id = doc.root;
let child_ids: Vec<_> = doc.children(&root_id).to_vec();
let traversal = AgentTraversal::new(doc);
if child_ids.is_empty() {
return; }
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
traversal.navigate_to(&session_id, root_id).unwrap();
let child_id = child_ids[0];
traversal.navigate_to(&session_id, child_id).unwrap();
let result = traversal.go_back(&session_id, 1).unwrap();
assert_eq!(result.position, root_id);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_expand_down() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.expand(
&session_id,
root_id,
ExpandDirection::Down,
ExpandOptions::new().with_depth(2),
);
assert!(result.is_ok());
let expansion = result.unwrap();
assert_eq!(expansion.root, root_id);
assert!(expansion.total_blocks >= 1);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_expand_with_depth_limit() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let config = SessionConfig::new().with_limits(SessionLimits {
max_expand_depth: 1,
..Default::default()
});
let session_id = traversal.create_session(config).unwrap();
let result = traversal.expand(
&session_id,
root_id,
ExpandDirection::Down,
ExpandOptions::new().with_depth(1),
);
assert!(result.is_ok());
let result = traversal.expand(
&session_id,
root_id,
ExpandDirection::Down,
ExpandOptions::new().with_depth(2),
);
assert!(matches!(
result,
Err(AgentError::DepthLimitExceeded { current: 2, max: 1 })
));
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_expand_up() {
let doc = create_test_document();
let root_id = doc.root;
let child_ids: Vec<_> = doc.children(&root_id).to_vec();
let traversal = AgentTraversal::new(doc);
if child_ids.is_empty() {
return; }
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.expand(
&session_id,
child_ids[0],
ExpandDirection::Up,
ExpandOptions::new().with_depth(2),
);
assert!(result.is_ok());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_by_role() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_by_pattern(&session_id, Some("paragraph"), None, None, None);
assert!(result.is_ok());
let find_result = result.unwrap();
assert!(find_result.total_searched > 0);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_by_tag() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_by_pattern(&session_id, None, Some("important"), None, None);
assert!(result.is_ok());
let find_result = result.unwrap();
assert!(find_result.total_searched > 0);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_by_pattern() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_by_pattern(&session_id, None, None, None, Some("introduction"));
assert!(result.is_ok());
let find_result = result.unwrap();
assert!(find_result.total_searched > 0);
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_search_without_rag_provider() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal
.search(&session_id, "test query", SearchOptions::new())
.await;
assert!(matches!(result, Err(AgentError::RagNotConfigured)));
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_search_with_mock_rag() {
let doc = create_test_document();
let root_id = doc.root;
let mut mock_rag = MockRagProvider::new();
mock_rag.add_result(root_id, 0.95, Some("Test content"));
let traversal =
AgentTraversal::new(doc).with_rag_provider(Arc::new(mock_rag) as Arc<dyn RagProvider>);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal
.search(
&session_id,
"test query",
SearchOptions::new().with_limit(10),
)
.await;
assert!(result.is_ok());
let search_result = result.unwrap();
assert!(!search_result.matches.is_empty());
assert_eq!(search_result.matches[0].block_id, root_id);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_view_block_full() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.view_block(&session_id, root_id, ViewMode::Full);
assert!(result.is_ok());
let view = result.unwrap();
assert_eq!(view.block_id, root_id);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_view_block_preview() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.view_block(&session_id, root_id, ViewMode::Preview { length: 10 });
assert!(result.is_ok());
let view = result.unwrap();
if let Some(content) = &view.content {
assert!(content.len() <= 10);
}
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_view_block_metadata() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.view_block(&session_id, root_id, ViewMode::Metadata);
assert!(result.is_ok());
let view = result.unwrap();
assert!(view.content.is_none());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_view_neighborhood() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.view_neighborhood(&session_id);
assert!(result.is_ok());
let view = result.unwrap();
assert!(!view.position.to_string().is_empty());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_path_same_node() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_path(&session_id, root_id, root_id, None);
assert!(result.is_ok());
let path = result.unwrap();
assert_eq!(path.len(), 1);
assert_eq!(path[0], root_id);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_path_root_to_child() {
let doc = create_test_document();
let root_id = doc.root;
let child_ids: Vec<_> = doc.children(&root_id).to_vec();
let traversal = AgentTraversal::new(doc);
if child_ids.is_empty() {
return; }
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_path(&session_id, root_id, child_ids[0], None);
assert!(result.is_ok());
let path = result.unwrap();
assert!(path.len() >= 2);
assert_eq!(path[0], root_id);
assert_eq!(*path.last().unwrap(), child_ids[0]);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_find_path_no_path_exists() {
let doc = Document::create();
let fake_id = fake_block_id();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.find_path(&session_id, BlockId::root(), fake_id, Some(5));
assert!(result.is_err());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_context_add_and_remove() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.context_add(
&session_id,
root_id,
Some("test reason".to_string()),
Some(0.9),
);
assert!(result.is_ok());
let result = traversal.context_remove(&session_id, root_id);
assert!(result.is_ok());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_context_clear() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
traversal
.context_add(&session_id, root_id, None, None)
.unwrap();
let result = traversal.context_clear(&session_id);
assert!(result.is_ok());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_context_focus() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.context_focus(&session_id, Some(root_id));
assert!(result.is_ok());
let result = traversal.context_focus(&session_id, None);
assert!(result.is_ok());
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_codegraph_context_session_operations_are_stateful() {
let doc = create_codegraph_document();
let traversal = AgentTraversal::new(doc.clone());
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let overview = traversal.codegraph_seed_overview(&session_id).unwrap();
assert!(overview.focus.is_some());
let file_id = ucp_codegraph::resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
traversal
.codegraph_expand_file(&session_id, file_id)
.unwrap();
let add_id = ucp_codegraph::resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let util_id =
ucp_codegraph::resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
traversal
.codegraph_expand_dependencies(&session_id, add_id, Some("uses_symbol"))
.unwrap();
traversal
.codegraph_expand_dependents(&session_id, util_id, Some("uses_symbol"))
.unwrap();
traversal
.codegraph_hydrate_source(&session_id, add_id, 1)
.unwrap();
let rendered = traversal
.render_codegraph_context(&session_id, ucp_codegraph::CodeGraphRenderConfig::default())
.unwrap();
assert!(rendered.contains("CodeGraph working set"));
assert!(rendered.contains("uses_symbol"));
let sessions = traversal.get_session(&session_id).unwrap();
let session = sessions.get(&session_id).unwrap();
assert!(session.context_blocks.contains(&file_id));
assert!(session.context_blocks.contains(&add_id));
assert!(session.context_blocks.contains(&util_id));
assert!(session.codegraph_context.is_some());
}
#[test]
fn test_context_add_results_without_search() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = traversal.context_add_results(&session_id);
assert!(matches!(result, Err(AgentError::NoResultsAvailable)));
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_context_add_results_after_find() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let find_result = traversal
.find_by_pattern(&session_id, None, None, None, Some(".*"))
.unwrap();
if !find_result.matches.is_empty() {
let result = traversal.context_add_results(&session_id);
assert!(result.is_ok());
let added = result.unwrap();
assert_eq!(added.len(), find_result.matches.len());
}
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_execute_ucl_goto() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let ucl_input = format!("GOTO {}", root_id);
let result = ucp_agent::execute_ucl(&traversal, &session_id, &ucl_input).await;
assert!(result.is_ok());
let results = result.unwrap();
assert_eq!(results.len(), 1);
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_execute_ucl_back_empty() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = ucp_agent::execute_ucl(&traversal, &session_id, "BACK").await;
assert!(result.is_err());
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_execute_ucl_find() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = ucp_agent::execute_ucl(&traversal, &session_id, "FIND ROLE=paragraph").await;
if let Err(ref e) = result {
eprintln!("Error: {:?}", e);
}
assert!(
result.is_ok(),
"UCL FIND should succeed: {:?}",
result.err()
);
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_execute_ucl_ctx_clear() {
let doc = create_test_document();
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let result = ucp_agent::execute_ucl(&traversal, &session_id, "CTX CLEAR").await;
assert!(result.is_ok());
traversal.close_session(&session_id).unwrap();
}
#[tokio::test]
async fn test_execute_ucl_multiple_commands() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let ucl_input = format!("GOTO {}\nFIND ROLE=paragraph\nCTX CLEAR", root_id);
let result = ucp_agent::execute_ucl(&traversal, &session_id, &ucl_input).await;
if let Err(ref e) = result {
eprintln!("Error: {:?}", e);
}
assert!(
result.is_ok(),
"Multiple UCL commands should succeed: {:?}",
result.err()
);
let results = result.unwrap();
assert_eq!(results.len(), 3);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_view_mode_preview() {
let mode = ViewMode::preview(50);
match mode {
ViewMode::Preview { length } => assert_eq!(length, 50),
_ => panic!("Expected Preview mode"),
}
}
#[test]
fn test_view_mode_adaptive() {
let mode = ViewMode::adaptive(0.7);
match mode {
ViewMode::Adaptive { interest_threshold } => {
assert!((interest_threshold - 0.7).abs() < 0.001)
}
_ => panic!("Expected Adaptive mode"),
}
}
#[test]
fn test_metrics_recorded_on_navigation() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
let _ = traversal.navigate_to(&session_id, root_id);
traversal.close_session(&session_id).unwrap();
}
#[test]
fn test_session_closed_after_completion() {
let doc = create_test_document();
let root_id = doc.root;
let traversal = AgentTraversal::new(doc);
let session_id = traversal.create_session(SessionConfig::default()).unwrap();
traversal.close_session(&session_id).unwrap();
let result = traversal.navigate_to(&session_id, root_id);
assert!(matches!(result, Err(AgentError::SessionNotFound(_))));
}
#[test]
fn test_multiple_sessions_independent() {
let doc = create_test_document();
let root_id = doc.root;
let child_ids: Vec<_> = doc.children(&root_id).to_vec();
let traversal = AgentTraversal::new(doc);
if child_ids.is_empty() {
return;
}
let session1 = traversal.create_session(SessionConfig::default()).unwrap();
let session2 = traversal.create_session(SessionConfig::default()).unwrap();
traversal.navigate_to(&session1, root_id).unwrap();
traversal.navigate_to(&session2, child_ids[0]).unwrap();
traversal.close_session(&session1).unwrap();
traversal.close_session(&session2).unwrap();
}