feat: enhance agentic workflow execution and improve response parsing

Author: jruokola

crates/codegraph-mcp/Cargo.toml

@@ -93,7 +93,7 @@ embeddings-jina = ["embeddings", "codegraph-vector/jina"]
 cloud = ["embeddings-jina", "codegraph-graph/surrealdb"]
 server-http = ["dep:axum", "dep:hyper", "dep:tower", "dep:http-body-util"]
 qwen-integration = []
-ai-enhanced = ["dep:codegraph-ai", "embeddings"]
+ai-enhanced = ["dep:codegraph-ai", "embeddings", "autoagents-experimental"]
 ai-enhanced-faiss = ["ai-enhanced", "faiss"]
 autoagents-experimental = ["dep:autoagents", "dep:autoagents-derive"]
 legacy-mcp-server = []

crates/codegraph-mcp/src/agentic_orchestrator.rs

@@ -512,6 +512,13 @@ impl AgenticOrchestrator {
 
     /// Parse LLM response to extract reasoning step
     fn parse_llm_response(&self, step_number: usize, response: &str) -> Result<ReasoningStep> {
+        Self::parse_llm_response_internal(step_number, response)
+    }
+
+    fn parse_llm_response_internal(
+        step_number: usize,
+        response: &str,
+    ) -> Result<ReasoningStep> {
         // Try to parse as JSON first
         if let Ok(parsed) = serde_json::from_str::<JsonValue>(response) {
             let reasoning = parsed["reasoning"]
@@ -523,8 +530,14 @@ impl AgenticOrchestrator {
 
             let (tool_name, tool_params) = if let Some(tool_call) = parsed["tool_call"].as_object()
             {
-                let name = tool_call.get("tool_name").and_then(|v| v.as_str());
-                let params = tool_call.get("parameters").cloned();
+                let name = tool_call
+                    .get("tool_name")
+                    .or_else(|| tool_call.get("name"))
+                    .and_then(|v| v.as_str());
+                let params = tool_call
+                    .get("parameters")
+                    .or_else(|| tool_call.get("arguments"))
+                    .cloned();
                 (name.map(|s| s.to_string()), params)
             } else {
                 (None, None)
@@ -580,8 +593,6 @@ mod tests {
 
     #[test]
     fn test_parse_valid_json_response() {
-        let orchestrator = create_test_orchestrator();
-
         let json_response = r#"{
             "reasoning": "I need to analyze dependencies",
             "tool_call": {
@@ -595,8 +606,7 @@ mod tests {
             "is_final": false
         }"#;
 
-        let step = orchestrator
-            .parse_llm_response(1, json_response)
+        let step = AgenticOrchestrator::parse_llm_response_internal(1, json_response)
             .expect("Should parse valid JSON");
 
         assert_eq!(step.step_number, 1);
@@ -610,37 +620,44 @@ mod tests {
 
     #[test]
     fn test_parse_final_response() {
-        let orchestrator = create_test_orchestrator();
-
         let json_response = r#"{
             "reasoning": "Based on the analysis, the answer is...",
             "tool_call": null,
             "is_final": true
         }"#;
 
-        let step = orchestrator
-            .parse_llm_response(1, json_response)
+        let step = AgenticOrchestrator::parse_llm_response_internal(1, json_response)
             .expect("Should parse final response");
 
         assert_eq!(step.step_number, 1);
         assert!(step.is_final);
         assert!(step.tool_name.is_none());
     }
 
-    // Helper to create test orchestrator
-    fn create_test_orchestrator() -> AgenticOrchestrator {
-        use codegraph_graph::GraphFunctions;
-        use std::sync::Arc;
-
-        // This is a minimal test setup - in real tests, we'd use mock objects
-        let graph_functions =
-            Arc::new(GraphFunctions::new_memory().expect("Failed to create test graph"));
-        let tool_executor = Arc::new(GraphToolExecutor::new(graph_functions));
+    #[test]
+    fn test_parse_tool_call_with_name_arguments_fields() {
+        let json_response = r#"{
+            "reasoning": "Need hub nodes",
+            "tool_call": {
+                "name": "get_hub_nodes",
+                "arguments": {
+                    "min_degree": 7
+                }
+            },
+            "is_final": false
+        }"#;
 
-        // For testing parse logic, we don't need a real LLM provider
-        // In real tests, we'd use a mock LLM provider
-        let llm_provider: Arc<dyn LLMProvider> = unimplemented!("Use mock in real tests");
+        let step = AgenticOrchestrator::parse_llm_response_internal(2, json_response)
+            .expect("Should parse alias fields");
 
-        AgenticOrchestrator::new(llm_provider, tool_executor, ContextTier::Medium)
+        assert_eq!(step.step_number, 2);
+        assert_eq!(step.tool_name, Some("get_hub_nodes".to_string()));
+        assert_eq!(
+            step.tool_params
+                .and_then(|p| p["min_degree"].as_i64())
+                .unwrap_or_default(),
+            7
+        );
+        assert!(!step.is_final);
     }
 }

crates/codegraph-mcp/src/autoagents/agent_builder.rs

@@ -10,7 +10,9 @@ use autoagents::llm::embedding::EmbeddingProvider;
 use autoagents::llm::error::LLMError;
 use autoagents::llm::models::{ModelListRequest, ModelListResponse, ModelsProvider};
 use autoagents::llm::{FunctionCall, ToolCall};
-use codegraph_ai::llm_provider::{LLMProvider as CodeGraphLLM, Message, MessageRole};
+use codegraph_ai::llm_provider::{
+    LLMProvider as CodeGraphLLM, LLMResponse, Message, MessageRole, ProviderCharacteristics,
+};
 use serde::Deserialize;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
@@ -220,7 +222,9 @@ struct CodeGraphLLMResponse {
 
 #[derive(Debug, Deserialize)]
 struct CodeGraphToolCall {
+    #[serde(alias = "name")]
     tool_name: String,
+    #[serde(alias = "arguments")]
     parameters: serde_json::Value,
 }
 
@@ -410,7 +414,9 @@ use crate::autoagents::codegraph_agent::CodeGraphAgentOutput;
 use autoagents::core::agent::memory::SlidingWindowMemory;
 use autoagents::core::agent::prebuilt::executor::ReActAgent;
 use autoagents::core::agent::AgentBuilder;
-use autoagents::core::agent::{AgentDeriveT, AgentExecutor, AgentHooks, AgentOutputT, Context, DirectAgentHandle, ExecutorConfig};
+use autoagents::core::agent::{
+    AgentDeriveT, AgentExecutor, AgentHooks, Context, DirectAgentHandle, ExecutorConfig,
+};
 use autoagents::core::error::Error as AutoAgentsError;
 use autoagents::core::tool::{shared_tools_to_boxes, ToolT};
 
@@ -608,6 +614,29 @@ mod tests {
         assert_eq!(aa_messages[1].role, ChatRole::User);
     }
 
+    #[test]
+    fn test_tool_calls_accepts_name_arguments_fields() {
+        let response = CodeGraphChatResponse {
+            content: r#"{
+                "reasoning": "Plan",
+                "tool_call": {
+                    "name": "get_hub_nodes",
+                    "arguments": {
+                        "min_degree": 4
+                    }
+                },
+                "is_final": false
+            }"#
+            .to_string(),
+            _total_tokens: 0,
+        };
+
+        let tool_calls = response.tool_calls().expect("tool call not parsed");
+        assert_eq!(tool_calls.len(), 1);
+        assert_eq!(tool_calls[0].function.name, "get_hub_nodes");
+        assert_eq!(tool_calls[0].function.arguments, "{\"min_degree\":4}");
+    }
+
     // Integration test for ChatProvider
     struct MockCodeGraphLLM;
 
@@ -617,12 +646,39 @@ mod tests {
             &self,
             messages: &[Message],
             _config: &codegraph_ai::llm_provider::GenerationConfig,
-        ) -> codegraph_ai::llm_provider::LLMResult<codegraph_ai::llm_provider::Response> {
-            Ok(codegraph_ai::llm_provider::Response {
+        ) -> codegraph_ai::llm_provider::LLMResult<LLMResponse> {
+            Ok(LLMResponse {
                 content: format!("Echo: {}", messages.last().unwrap().content),
-                total_tokens: 10,
+                total_tokens: Some(10),
+                prompt_tokens: None,
+                completion_tokens: None,
+                finish_reason: Some("stop".to_string()),
+                model: "mock".to_string(),
             })
         }
+
+        async fn is_available(&self) -> bool {
+            true
+        }
+
+        fn provider_name(&self) -> &str {
+            "mock"
+        }
+
+        fn model_name(&self) -> &str {
+            "mock-model"
+        }
+
+        fn characteristics(&self) -> ProviderCharacteristics {
+            ProviderCharacteristics {
+                max_tokens: 4096,
+                avg_latency_ms: 1,
+                rpm_limit: None,
+                tpm_limit: None,
+                supports_streaming: false,
+                supports_functions: true,
+            }
+        }
     }
 
     #[tokio::test]

crates/codegraph-mcp/src/heartbeat.rs

@@ -231,7 +231,7 @@ impl HeartbeatManager {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use tokio::time::{advance, pause};
+    use tokio::time::sleep;
 
     #[tokio::test]
     async fn test_heartbeat_config_default() {
@@ -266,8 +266,6 @@ mod tests {
 
     #[tokio::test]
     async fn test_pong_handling() {
-        pause();
-
         let config = HeartbeatConfig {
             interval: Duration::from_millis(100),
             timeout: Duration::from_millis(50),
@@ -278,7 +276,7 @@ mod tests {
         monitor.on_pong_received(1).await;
         assert_eq!(monitor.state().await, HeartbeatState::Healthy);
 
-        advance(Duration::from_millis(200)).await;
+        sleep(Duration::from_millis(200)).await;
 
         monitor.on_pong_received(100).await;
         assert_eq!(monitor.state().await, HeartbeatState::Healthy);