Merge pull request #8 from sopaco/v2

V2

Author: sopaco

Cargo.lock

@@ -18,6 +18,7 @@ adk-tool = "0.2.1"
 # Async runtime
 tokio = { workspace = true }
 async-trait = "0.1"
+async-stream = "0.3"
 futures = { workspace = true }
 
 # Error handling

crates/cowork-core/Cargo.toml

@@ -1,27 +1,83 @@
 use adk_core::{Agent, Event, AdkError, InvocationContext};
 use async_trait::async_trait;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::pin::Pin;
 use std::task::{Context as TaskContext, Poll};
 use futures::{Stream, Future};
 use dialoguer::{Select, Input, theme::ColorfulTheme};
 
+type AgentOutput = Pin<Box<dyn Stream<Item = Result<Event, AdkError>> + Send>>;
+
 pub struct ResilientAgent {
     inner: Arc<dyn Agent>,
     subs: Vec<Arc<dyn Agent>>,
+    retry_count: Arc<AtomicU32>,
 }
 
 impl ResilientAgent {
+    const MAX_RETRY_ATTEMPTS: u32 = 3;
+    
     pub fn new(inner: Arc<dyn Agent>) -> Self {
         Self {
             inner: inner.clone(),
             subs: vec![inner],
+            retry_count: Arc::new(AtomicU32::new(0)),
         }
     }
+    
+    // Helper for immediate errors (recursion in async fn)
+    async fn handle_error(&self, context: Arc<dyn InvocationContext>, e: AdkError) -> Result<AgentOutput, AdkError> {
+        let current_retry = self.retry_count.fetch_add(1, Ordering::SeqCst);
+        
+        // Check if max retry attempts reached
+        if current_retry >= Self::MAX_RETRY_ATTEMPTS {
+            println!("\n❌ Maximum retry attempts ({}) reached.", Self::MAX_RETRY_ATTEMPTS);
+            self.retry_count.store(0, Ordering::SeqCst); // Reset counter
+            return Err(AdkError::Tool(format!(
+                "Agent '{}' failed after {} retry attempts", 
+                self.name(), 
+                Self::MAX_RETRY_ATTEMPTS
+            )));
+        }
+        
+        println!("\n⚠️  Agent '{}' encountered error: {}", self.name(), e);
+        println!("The agent loop limit has been exceeded.");
+        println!("Retry attempt {}/{}", current_retry + 1, Self::MAX_RETRY_ATTEMPTS);
+         
+        let selections = &["Retry (reset counter)", "Provide Guidance & Retry", "Abort"];
+        let selection = Select::with_theme(&ColorfulTheme::default())
+           .with_prompt("How would you like to proceed?")
+           .default(0)
+           .items(&selections[..])
+           .interact()
+           .unwrap_or(2);
+
+        match selection {
+           0 => {
+               println!("🔄 Retrying agent execution...");
+               return self.run(context).await;
+           },
+           1 => {
+               let input: String = Input::with_theme(&ColorfulTheme::default())
+                   .with_prompt("Please provide guidance for the agent")
+                   .interact_text()
+                   .unwrap_or_default();
+               
+               if !input.is_empty() {
+                   println!("(Note: User guidance provided: '{}' - but context injection is not implemented. Retrying anyway.)", input);
+               }
+               println!("🔄 Retrying with new guidance...");
+               return self.run(context).await;
+           },
+           _ => {
+               self.retry_count.store(0, Ordering::SeqCst); // Reset counter on abort
+               return Err(e);
+           }
+        }
+   }
 }
 
-type AgentOutput = Pin<Box<dyn Stream<Item = Result<Event, AdkError>> + Send>>;
-
 #[async_trait]
 impl Agent for ResilientAgent {
     fn name(&self) -> &str {
@@ -40,11 +96,14 @@ impl Agent for ResilientAgent {
         // Initial run
         match self.inner.run(context.clone()).await {
             Ok(stream) => {
+                // Success - reset retry counter
+                self.retry_count.store(0, Ordering::SeqCst);
                 // Wrap the stream to handle errors during iteration
                 Ok(Box::pin(ResilientStream::new(
                     self.inner.clone(),
                     context,
-                    stream
+                    stream,
+                    self.retry_count.clone(),
                 )))
             },
             Err(e) => {
@@ -61,42 +120,6 @@ impl Agent for ResilientAgent {
     }
 }
 
-impl ResilientAgent {
-    // Helper for immediate errors (recursion in async fn)
-    async fn handle_error(&self, context: Arc<dyn InvocationContext>, e: AdkError) -> Result<AgentOutput, AdkError> {
-         println!("\n⚠️  Agent '{}' encountered error: {}", self.name(), e);
-         println!("The agent loop limit has been exceeded.");
-         
-         let selections = &["Retry (reset counter)", "Provide Guidance & Retry", "Abort"];
-         let selection = Select::with_theme(&ColorfulTheme::default())
-            .with_prompt("How would you like to proceed?")
-            .default(0)
-            .items(&selections[..])
-            .interact()
-            .unwrap_or(2);
-
-         match selection {
-            0 => {
-                println!("🔄 Retrying agent execution...");
-                return self.run(context).await;
-            },
-            1 => {
-                let input: String = Input::with_theme(&ColorfulTheme::default())
-                    .with_prompt("Please provide guidance for the agent")
-                    .interact_text()
-                    .unwrap_or_default();
-                
-                if !input.is_empty() {
-                    println!("(Note: User guidance provided: '{}' - but context injection is not implemented. Retrying anyway.)", input);
-                }
-                println!("🔄 Retrying with new guidance...");
-                return self.run(context).await;
-            },
-            _ => return Err(e),
-         }
-    }
-}
-
 // ============================================================================
 // ResilientStream Implementation
 // ============================================================================
@@ -111,20 +134,23 @@ struct ResilientStream {
     context: Arc<dyn InvocationContext>,
     state: StreamState,
     agent_name: String, // Cached for logging
+    retry_count: Arc<AtomicU32>,
 }
 
 impl ResilientStream {
     fn new(
         inner_agent: Arc<dyn Agent>,
         context: Arc<dyn InvocationContext>,
         stream: AgentOutput,
+        retry_count: Arc<AtomicU32>,
     ) -> Self {
         let agent_name = inner_agent.name().to_string();
         Self {
             inner_agent,
             context,
             state: StreamState::Streaming(stream),
             agent_name,
+            retry_count,
         }
     }
 

crates/cowork-core/src/agents/hitl.rs

@@ -1,18 +1,17 @@
 // Agents module - Agent builders using adk-rust
 // 
-// IMPORTANT: This file solves a CRITICAL bug where SequentialAgent stops after
-// the first LoopAgent completes. 
+// IMPORTANT: This file uses StageExecutor instead of SequentialAgent to allow
+// LoopAgents to use ExitLoopTool without affecting other stages.
 //
-// PROBLEM: When a sub-agent in LoopAgent calls exit_loop(), it terminates the
-// ENTIRE SequentialAgent, not just the LoopAgent. This is adk-rust's design.
-//
-// SOLUTION: Remove exit_loop tools and use max_iterations=1 to let LoopAgent
-// complete naturally, allowing SequentialAgent to continue to next agent.
+// SOLUTION: StageExecutor isolates each stage's escalate flag, so when a
+// sub-agent in LoopAgent calls exit_loop(), it only terminates that specific
+// LoopAgent, not the entire workflow.
 
 use crate::instructions::*;
 use crate::tools::*;
 use adk_agent::{LlmAgentBuilder, LoopAgent};
 use adk_core::{Llm, IncludeContents};
+use adk_tool::ExitLoopTool;
 use anyhow::Result;
 use std::sync::Arc;
 
@@ -60,15 +59,14 @@ pub fn create_prd_loop(model: Arc<dyn Llm>, session_id: &str) -> Result<Arc<dyn
         .model(model)
         .tool(Arc::new(ReadFileTool))
         .tool(Arc::new(GetRequirementsTool::new(session.clone())))
-        .tool(Arc::new(ProvideFeedbackTool::new(session.clone())))
-        .include_contents(IncludeContents::None)
+        .tool(Arc::new(ProvideFeedbackTool::new(session.clone()))) // Write feedback to file when checks fail
+        .tool(Arc::new(ExitLoopTool::new())) // Exit loop when checks pass
+        .tool(Arc::new(RequestHumanReviewTool::new(session.clone())))
+        .include_contents(IncludeContents::Default)
         .build()?;
 
-    let mut loop_agent = LoopAgent::new(
-        "prd_loop",
-        vec![Arc::new(prd_actor), Arc::new(prd_critic)],
-    );
-    loop_agent = loop_agent.with_max_iterations(3); // Allow up to 3 attempts for Actor to fix issues
+    let mut loop_agent = LoopAgent::new("prd_loop", vec![Arc::new(prd_actor), Arc::new(prd_critic)]);
+    loop_agent = loop_agent.with_max_iterations(3); // Loop will complete naturally after 3 iterations
 
     Ok(Arc::new(ResilientAgent::new(Arc::new(loop_agent))))
 }
@@ -99,11 +97,13 @@ pub fn create_design_loop(model: Arc<dyn Llm>, session_id: &str) -> Result<Arc<d
         .tool(Arc::new(GetDesignTool::new(session.clone())))
         .tool(Arc::new(CheckFeatureCoverageTool::new(session.clone())))
         .tool(Arc::new(ProvideFeedbackTool::new(session.clone())))
-        .include_contents(IncludeContents::None)
+        .tool(Arc::new(ExitLoopTool::new()))
+        .tool(Arc::new(RequestHumanReviewTool::new(session.clone())))
+        .include_contents(IncludeContents::Default)
         .build()?;
 
     let mut loop_agent = LoopAgent::new("design_loop", vec![Arc::new(design_actor), Arc::new(design_critic)]);
-    loop_agent = loop_agent.with_max_iterations(3); // Allow up to 3 attempts
+    loop_agent = loop_agent.with_max_iterations(3); // Loop will complete naturally after 3 iterations
 
     Ok(Arc::new(ResilientAgent::new(Arc::new(loop_agent))))
 }
@@ -137,11 +137,13 @@ pub fn create_plan_loop(model: Arc<dyn Llm>, session_id: &str) -> Result<Arc<dyn
         .tool(Arc::new(GetDesignTool::new(session.clone())))
         .tool(Arc::new(CheckTaskDependenciesTool::new(session.clone())))
         .tool(Arc::new(ProvideFeedbackTool::new(session.clone())))
-        .include_contents(IncludeContents::None)
+        .tool(Arc::new(ExitLoopTool::new()))
+        .tool(Arc::new(RequestHumanReviewTool::new(session.clone())))
+        .include_contents(IncludeContents::Default)
         .build()?;
 
     let mut loop_agent = LoopAgent::new("plan_loop", vec![Arc::new(plan_actor), Arc::new(plan_critic)]);
-    loop_agent = loop_agent.with_max_iterations(3); // Allow up to 3 attempts
+    loop_agent = loop_agent.with_max_iterations(3); // Loop will complete naturally after 3 iterations
 
     Ok(Arc::new(ResilientAgent::new(Arc::new(loop_agent))))
 }
@@ -157,9 +159,10 @@ pub fn create_coding_loop(model: Arc<dyn Llm>, session_id: &str) -> Result<Arc<d
         .instruction(CODING_ACTOR_INSTRUCTION)
         .model(model.clone())
         .tool(Arc::new(GetPlanTool::new(session.clone())))
+        .tool(Arc::new(ReviewWithFeedbackContentTool)) // Read Critic feedback
         .tool(Arc::new(UpdateTaskStatusTool::new(session.clone())))
         .tool(Arc::new(UpdateFeatureStatusTool::new(session.clone())))
-        // Task management tools - NEW
+        // Task management tools
         .tool(Arc::new(CreateTaskTool::new(session.clone())))
         .tool(Arc::new(UpdateTaskTool::new(session.clone())))
         .tool(Arc::new(DeleteTaskTool::new(session.clone())))
@@ -180,13 +183,13 @@ pub fn create_coding_loop(model: Arc<dyn Llm>, session_id: &str) -> Result<Arc<d
         .tool(Arc::new(ListFilesTool))
         .tool(Arc::new(RunCommandTool))
         .tool(Arc::new(ProvideFeedbackTool::new(session.clone())))
-        // Replanning request - NEW
+        .tool(Arc::new(ExitLoopTool::new()))
         .tool(Arc::new(RequestReplanningTool::new(session.clone())))
-        .include_contents(IncludeContents::None)
+        .include_contents(IncludeContents::Default)
         .build()?;
 
     let mut loop_agent = LoopAgent::new("coding_loop", vec![Arc::new(coding_actor), Arc::new(coding_critic)]);
-    loop_agent = loop_agent.with_max_iterations(5);
+    loop_agent = loop_agent.with_max_iterations(5); // Loop will complete naturally after 5 iterations
 
     Ok(Arc::new(ResilientAgent::new(Arc::new(loop_agent))))
 }

crates/cowork-core/src/agents/mod.rs

@@ -54,16 +54,30 @@ During implementation, you may discover that the plan needs adjustments. You now
 - **Stay focused**: Don't over-plan; focus on what's needed for current implementation
 - **Maintain consistency**: Keep task IDs, dependencies, and status aligned
 
+## Handle Critic Feedback (IF IN ITERATION 2+):
+**IMPORTANT**: In iterations after the first one, check the conversation history for Critic's feedback:
+
+1. **Look at the previous messages** - Critic's feedback is in the conversation history
+2. **If Critic said code is incomplete or has issues**:
+   - Read exactly what issues were mentioned
+   - Complete any missing tasks
+   - Fix any code quality issues
+   - Simplify over-engineered code if needed
+3. **If Critic requested replanning**: Acknowledge (human will review)
+4. **If no issues mentioned** - Critic approved and you're done!
+
+**Remember**: You can SEE Critic's messages in the conversation. Read them and take action.
+
 # Tools
 - get_plan()
 - read_file(path)
 - write_file(path, content)
 - list_files(path)
 - update_task_status(task_id, status)
 - update_feature_status(feature_id, status)
-- create_task(title, description, feature_id, component_id, files_to_create, dependencies, acceptance_criteria) ← NEW
-- update_task(task_id, reason, title?, description?, dependencies?, files_to_create?, acceptance_criteria?) ← NEW
-- delete_task(task_id, reason) ← NEW
+- create_task(title, description, feature_id, component_id, files_to_create, dependencies, acceptance_criteria)
+- update_task(task_id, reason, title?, description?, dependencies?, files_to_create?, acceptance_criteria?)
+- delete_task(task_id, reason)
 
 # Code Style - SIMPLE APPROACH
 ```
@@ -152,13 +166,19 @@ The request will be recorded and reviewed by the Check Agent, which can trigger
 # Exit Condition
 - When ALL tasks show status="completed" AND key files exist, approve immediately and stop
 
+## Decision Flow:
+1. If major architectural issues found → call `request_replanning()`
+2. If minor issues or incomplete work → call `provide_feedback(feedback_type="incomplete", severity="medium", details="<what's pending>", suggested_fix="<complete these tasks>")`
+3. If everything is complete and good → call `exit_loop()`
+
 # Tools
 - get_plan()
 - read_file(path)
 - list_files(path)  ← Use this to verify files exist!
 - run_command(command)  ← Only for simple checks, not for tests/lint
-- provide_feedback(feedback_type, severity, details, suggested_fix)
-- request_replanning(issue_type, severity, details, affected_features, suggested_approach) ← NEW
+- provide_feedback(feedback_type, severity, details, suggested_fix) - Report incomplete work (Actor will read this)
+- exit_loop() - **MUST CALL** when all tasks completed and code is good (exits this loop only, other stages continue)
+- request_replanning(issue_type, severity, details, affected_features, suggested_approach) - Call for major issues
 
 # Example - All Tasks Done
 ```
@@ -168,8 +188,20 @@ The request will be recorded and reviewed by the Check Agent, which can trigger
 4. # Returns: ["index.html", "style.css", "script.js"] - files exist!
 5. read_file("index.html")
 6. # Looks good, simple HTML structure
-7. "✅ All 12 tasks completed. Files created: index.html, style.css, script.js. Code is simple and clear. Project ready!"
-8. STOP (no more iterations)
+7. exit_loop() - Exit the coding loop, workflow continues to check stage
+```
+
+# Example - Issues Found
+```
+1. get_plan()
+2. # Returns: 8 completed, 4 pending
+3. provide_feedback(
+     feedback_type="incomplete",
+     severity="medium",
+     details="4 tasks still pending: TASK-009 (Create login form), TASK-010 (Add validation), TASK-011 (Style buttons), TASK-012 (Add error messages)",
+     suggested_fix="Complete the 4 pending tasks before marking work as done"
+   )
+# Actor will read this feedback file and complete the tasks
 ```
 
 # Example - Tasks Complete but Files Missing