From 9d6d1746c66dc38f7ed81eeb3ba5750dd0c1b11a Mon Sep 17 00:00:00 2001
From: Ed_ <edwardgz@gmail.com>
Date: Thu, 26 Feb 2026 20:10:15 -0500
Subject: [PATCH] feat(mma): Implement context injection using ASTParser in
 run_worker_lifecycle

---
 multi_agent_conductor.py       | 24 +++++++++++++++-
 tests/test_conductor_engine.py | 52 ++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/multi_agent_conductor.py b/multi_agent_conductor.py
index 79f6c12..a8a257a 100644
--- a/multi_agent_conductor.py
+++ b/multi_agent_conductor.py
@@ -1,5 +1,7 @@
 import ai_client
+from typing import List, Optional
 from models import Ticket, Track, WorkerContext
+from file_cache import ASTParser
 
 class ConductorEngine:
     """
@@ -34,15 +36,35 @@ class ConductorEngine:
                 )
                 run_worker_lifecycle(ticket, context)
 
-def run_worker_lifecycle(ticket: Ticket, context: WorkerContext):
+def run_worker_lifecycle(ticket: Ticket, context: WorkerContext, context_files: List[str] = None):
     """
     Simulates the lifecycle of a single agent working on a ticket.
     Calls the AI client and updates the ticket status based on the response.
     """
+    context_injection = ""
+    if context_files:
+        parser = ASTParser(language="python")
+        for i, file_path in enumerate(context_files):
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                if i == 0:
+                    view = parser.get_curated_view(content)
+                else:
+                    view = parser.get_skeleton(content)
+                context_injection += f"\nFile: {file_path}\n{view}\n"
+            except Exception as e:
+                context_injection += f"\nError reading {file_path}: {e}\n"
+
     # Build a prompt for the worker
     user_message = (
         f"You are assigned to Ticket {ticket.id}.\n"
         f"Task Description: {ticket.description}\n"
+    )
+    if context_injection:
+        user_message += f"\nContext Files:\n{context_injection}\n"
+        
+    user_message += (
         "Please complete this task. If you are blocked and cannot proceed, "
         "start your response with 'BLOCKED' and explain why."
     )
diff --git a/tests/test_conductor_engine.py b/tests/test_conductor_engine.py
index c48d398..5294377 100644
--- a/tests/test_conductor_engine.py
+++ b/tests/test_conductor_engine.py
@@ -69,6 +69,58 @@ def test_run_worker_lifecycle_calls_ai_client_send():
         # user_message is passed as a keyword argument
         assert ticket.description in kwargs["user_message"]
 
+def test_run_worker_lifecycle_context_injection():
+    """
+    Test that run_worker_lifecycle can take a context_files list and injects AST views into the prompt.
+    """
+    ticket = Ticket(id="T1", description="Task 1", status="todo", assigned_to="worker1")
+    context = WorkerContext(ticket_id="T1", model_name="test-model", messages=[])
+    context_files = ["primary.py", "secondary.py"]
+    
+    from multi_agent_conductor import run_worker_lifecycle
+    
+    # We mock ASTParser which is expected to be imported in multi_agent_conductor
+    with patch("ai_client.send") as mock_send, \
+         patch("multi_agent_conductor.ASTParser") as mock_ast_parser_class, \
+         patch("builtins.open", new_callable=MagicMock) as mock_open:
+        
+        # Setup open mock to return different content for different files
+        file_contents = {
+            "primary.py": "def primary(): pass",
+            "secondary.py": "def secondary(): pass"
+        }
+        
+        def mock_open_side_effect(file, *args, **kwargs):
+            content = file_contents.get(file, "")
+            mock_file = MagicMock()
+            mock_file.read.return_value = content
+            mock_file.__enter__.return_value = mock_file
+            return mock_file
+            
+        mock_open.side_effect = mock_open_side_effect
+        
+        # Setup ASTParser mock
+        mock_ast_parser = mock_ast_parser_class.return_value
+        mock_ast_parser.get_curated_view.return_value = "CURATED VIEW"
+        mock_ast_parser.get_skeleton.return_value = "SKELETON VIEW"
+        
+        mock_send.return_value = "Success"
+        
+        run_worker_lifecycle(ticket, context, context_files=context_files)
+        
+        # Verify ASTParser calls: 
+        # First file (primary) should get curated view, others (secondary) get skeleton
+        mock_ast_parser.get_curated_view.assert_called_once_with("def primary(): pass")
+        mock_ast_parser.get_skeleton.assert_called_once_with("def secondary(): pass")
+        
+        # Verify user_message contains the views
+        _, kwargs = mock_send.call_args
+        user_message = kwargs["user_message"]
+        assert "CURATED VIEW" in user_message
+        assert "SKELETON VIEW" in user_message
+        assert "primary.py" in user_message
+        assert "secondary.py" in user_message
+
 def test_run_worker_lifecycle_handles_blocked_response():
     """
     Test that run_worker_lifecycle marks the ticket as blocked if the AI indicates it cannot proceed.