From a86ffd3f495799bb93376341964e6ae2057c10cb Mon Sep 17 00:00:00 2001 From: Adil Hafeez Date: Thu, 18 Dec 2025 13:43:45 -0800 Subject: [PATCH] fix tests and update tracing --- .../src/handlers/agent_chat_completions.rs | 39 ++++++++++++++++++- .../src/handlers/agent_selector.rs | 2 +- .../src/handlers/integration_tests.rs | 6 +-- .../src/handlers/pipeline_processor.rs | 2 +- crates/brightstaff/src/tracing/constants.rs | 5 ++- 5 files changed, 47 insertions(+), 7 deletions(-) diff --git a/crates/brightstaff/src/handlers/agent_chat_completions.rs b/crates/brightstaff/src/handlers/agent_chat_completions.rs index a0be2442..576a4619 100644 --- a/crates/brightstaff/src/handlers/agent_chat_completions.rs +++ b/crates/brightstaff/src/handlers/agent_chat_completions.rs @@ -216,10 +216,45 @@ async fn handle_agent_chat( }; // Select appropriate agents using arch orchestrator llm model + let selection_span_id = generate_random_span_id(); + let selection_start_time = SystemTime::now(); + let selection_start_instant = Instant::now(); + let selected_agents = agent_selector .select_agents(&message, &listener, trace_parent.clone()) .await?; + // Record agent selection span + let selection_end_time = SystemTime::now(); + let selection_elapsed = selection_start_instant.elapsed(); + let selection_operation_name = OperationNameBuilder::new() + .with_method("POST") + .with_path("/agents/select") + .with_target(&listener.name) + .build(); + + let mut selection_span_builder = SpanBuilder::new(&selection_operation_name) + .with_span_id(selection_span_id) + .with_kind(SpanKind::Internal) + .with_start_time(selection_start_time) + .with_end_time(selection_end_time) + .with_attribute(http::METHOD, "POST") + .with_attribute(http::TARGET, "/agents/select") + .with_attribute("selection.listener", listener.name.clone()) + .with_attribute("selection.agent_count", selected_agents.len().to_string()) + .with_attribute("selection.agents", selected_agents.iter().map(|a| a.id.as_str()).collect::>().join(",")) + .with_attribute("duration_ms", format!("{:.2}", selection_elapsed.as_secs_f64() * 1000.0)); + + if !trace_id.is_empty() { + selection_span_builder = selection_span_builder.with_trace_id(trace_id.clone()); + } + if let Some(parent_id) = parent_span_id.clone() { + selection_span_builder = selection_span_builder.with_parent_span_id(parent_id); + } + + let selection_span = selection_span_builder.build(); + trace_collector.record_span(operation_component::ORCHESTRATOR, selection_span); + info!("Selected {} agent(s) for execution", selected_agents.len()); // Execute agents sequentially, passing output from one to the next @@ -241,6 +276,9 @@ async fn handle_agent_chat( let agent_start_instant = Instant::now(); let span_id = generate_random_span_id(); + // Get agent name + let agent_name = selected_agent.id.clone(); + // Process the filter chain let chat_history = pipeline_processor .process_filter_chain( @@ -255,7 +293,6 @@ async fn handle_agent_chat( .await?; // Get agent details and invoke - let agent_name = selected_agent.id.clone(); let agent = agent_map.get(&agent_name).unwrap(); debug!("Invoking agent: {}", agent_name); diff --git a/crates/brightstaff/src/handlers/agent_selector.rs b/crates/brightstaff/src/handlers/agent_selector.rs index e08bd509..ad281d7c 100644 --- a/crates/brightstaff/src/handlers/agent_selector.rs +++ b/crates/brightstaff/src/handlers/agent_selector.rs @@ -188,7 +188,7 @@ mod tests { id: name.to_string(), description: Some(description.to_string()), default: Some(is_default), - filter_chain: vec![name.to_string()], + filter_chain: Some(vec![name.to_string()]), } } diff --git a/crates/brightstaff/src/handlers/integration_tests.rs b/crates/brightstaff/src/handlers/integration_tests.rs index 7944e448..2dae9808 100644 --- a/crates/brightstaff/src/handlers/integration_tests.rs +++ b/crates/brightstaff/src/handlers/integration_tests.rs @@ -64,7 +64,7 @@ mod integration_tests { let agent_pipeline = AgentFilterChain { id: "terminal-agent".to_string(), - filter_chain: vec!["filter-agent".to_string(), "terminal-agent".to_string()], + filter_chain: Some(vec!["filter-agent".to_string(), "terminal-agent".to_string()]), description: Some("Test pipeline".to_string()), default: Some(true), }; @@ -104,7 +104,7 @@ mod integration_tests { // Create a pipeline with empty filter chain to avoid network calls let test_pipeline = AgentFilterChain { id: "terminal-agent".to_string(), - filter_chain: vec![], // Empty filter chain - no network calls needed + filter_chain: Some(vec![]), // Empty filter chain - no network calls needed description: None, default: None, }; @@ -143,7 +143,7 @@ mod integration_tests { #[tokio::test] async fn test_error_handling_flow() { - let router_service = create_test_router_service(); + let router_service = create_test_orchestrator_service(); let agent_selector = AgentSelector::new(router_service); // Test listener not found diff --git a/crates/brightstaff/src/handlers/pipeline_processor.rs b/crates/brightstaff/src/handlers/pipeline_processor.rs index c5bd89ec..c0cd1cef 100644 --- a/crates/brightstaff/src/handlers/pipeline_processor.rs +++ b/crates/brightstaff/src/handlers/pipeline_processor.rs @@ -740,7 +740,7 @@ mod tests { fn create_test_pipeline(agents: Vec<&str>) -> AgentFilterChain { AgentFilterChain { id: "test-agent".to_string(), - filter_chain: agents.iter().map(|s| s.to_string()).collect(), + filter_chain: Some(agents.iter().map(|s| s.to_string()).collect()), description: None, default: None, } diff --git a/crates/brightstaff/src/tracing/constants.rs b/crates/brightstaff/src/tracing/constants.rs index d3eafb2e..1edc4d3d 100644 --- a/crates/brightstaff/src/tracing/constants.rs +++ b/crates/brightstaff/src/tracing/constants.rs @@ -150,9 +150,12 @@ pub mod operation_component { /// Inbound request handling pub const INBOUND: &str = "plano(inbound)"; - /// Routing decision phase + /// Orchestrator for llm route selection pub const ROUTING: &str = "plano(routing)"; + /// Orchestrator for agent selection + pub const ORCHESTRATOR: &str = "plano(orchestrator)"; + /// Handoff to upstream service pub const HANDOFF: &str = "plano(handoff)";