Allow individual tasks to be marked completed.

ra_aid/agent_utils.py

@@ -218,16 +218,6 @@ def run_agent_with_retry(agent, prompt: str, config: dict) -> Optional[str]:
                 config
             ):
                 print_agent_output(chunk)
-                
-                # Check for task completion after each chunk
-                if _global_memory.get('task_completed'):
-                    completion_msg = _global_memory.get('completion_message', 'Task was completed successfully.')
-                    console.print(Panel(
-                        Markdown(completion_msg),
-                        title="✅ Task Completed",
-                        style="green"
-                    ))
-                    return completion_msg
             break
         except (InternalServerError, APITimeoutError, RateLimitError, APIError) as e:
             if attempt == max_retries - 1:

ra_aid/tool_configs.py

@@ -6,7 +6,8 @@ from ra_aid.tools import (
     emit_key_snippets, delete_key_snippets, delete_tasks,
     request_implementation, read_file_tool,
     fuzzy_find_project_files, ripgrep_search, list_directory_tree,
-    swap_task_order, monorepo_detected, existing_project_detected, ui_detected
+    swap_task_order, monorepo_detected, existing_project_detected, ui_detected,
+    task_completed
 )
 from ra_aid.tools.memory import one_shot_completed
 from ra_aid.tools.agent import request_research, request_task_implementation
@@ -94,6 +95,9 @@ def get_implementation_tools(expert_enabled: bool = True) -> list:
     
     # Add modification tools since it's not research-only
     tools.extend(MODIFICATION_TOOLS)
+    tools.extend([
+        task_completed
+    ])
     
     # Add expert tools if enabled
     if expert_enabled:

ra_aid/tools/__init__.py

@@ -12,7 +12,7 @@ from .ripgrep import ripgrep_search
 from .memory import (
     delete_tasks, emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
     request_implementation, delete_key_facts,
-    emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order
+    emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order, task_completed
 )
 
 __all__ = [
@@ -41,5 +41,6 @@ __all__ = [
     'monorepo_detected',
     'existing_project_detected',
     'ui_detected',
-    'ask_human'
+    'ask_human',
+    'task_completed'
 ]

ra_aid/tools/memory.py

@@ -299,6 +299,21 @@ def one_shot_completed(message: str) -> str:
     _global_memory['completion_message'] = message
     return message
 
+@tool("task_completed")
+def task_completed(message: str) -> str:
+    """Mark the current task as completed with a completion message.
+    
+    Args:
+        message: Message explaining how/why the task is complete
+        
+    Returns:
+        The completion message
+    """
+    _global_memory['task_completed'] = True
+    _global_memory['completion_message'] = message
+    console.print(Panel(Markdown(message), title="✅ Task Completed"))
+    return "Completion noted."
+
 def get_related_files() -> Set[str]:
     """Get the current set of related files.