diagram-to-iac 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

diagram_to_iac/agents/supervisor_langgraph/agent.py

@@ -23,15 +23,20 @@ Architecture:
 
 from __future__ import annotations
 
+import fnmatch
+import json
 import logging
 import os
-import fnmatch
 import uuid
-from datetime import datetime
-from typing import Optional, Dict, List, Set, TypedDict, Annotated
-
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Dict, Any, Optional, List, Annotated
+from typing_extensions import TypedDict
+from dataclasses import dataclass, asdict, field
+from enum import Enum
 import yaml
-from pydantic import BaseModel, Field
+
+from pydantic import BaseModel, Field, validator
 from langchain_core.messages import HumanMessage, BaseMessage
 from langgraph.graph import StateGraph, END
 from langgraph.checkpoint.memory import MemorySaver
@@ -134,9 +139,8 @@ class SupervisorAgent(AgentBase):
         git_agent: Optional[GitAgent] = None,
         shell_agent: Optional[ShellAgent] = None,
         terraform_agent: Optional[TerraformAgent] = None,
-
+        registry: Optional["RunRegistry"] = None,
         demonstrator: Optional[DryRunDemonstrator] = None,
-
         issue_tracker: Optional[IssueTracker] = None,
 
     ) -> None:
@@ -207,6 +211,16 @@ class SupervisorAgent(AgentBase):
         self.checkpointer = MemorySaver()
         self.logger.info("MemorySaver checkpointer initialized")
 
+        # Initialize run registry for issue linking and metadata tracking
+        from diagram_to_iac.core.registry import get_default_registry
+        self.run_registry = registry or get_default_registry()
+        self.logger.info("Run registry initialized for issue tracking")
+
+        # Check for PR merge context
+        self.pr_merge_context = self._detect_pr_merge_context()
+        if self.pr_merge_context:
+            self.logger.info(f"PR merge context detected: PR #{self.pr_merge_context.get('pr_number')} -> SHA {self.pr_merge_context.get('merged_sha')}")
+
         # Issue tracker for deduplicating issues
         self.issue_tracker = issue_tracker or IssueTracker()
 
@@ -612,22 +626,24 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                 )
             )
 
-            if git_result.error_message:
+            if git_result.artifacts and git_result.artifacts.get('error_message'):
+                error_message = git_result.artifacts.get('error_message')
                 self.logger.error(
-                    f"Repository cloning failed: {git_result.error_message}"
+                    f"Repository cloning failed: {error_message}"
                 )
                 return {
-                    "final_result": f"Repository cloning failed: {git_result.error_message}",
-                    "error_message": git_result.error_message,
+                    "final_result": f"Repository cloning failed: {error_message}",
+                    "error_message": error_message,
                     "operation_type": "clone_error",
                 }
 
             # Update state with repo path and continue to stack detection
+            repo_path = git_result.artifacts.get('repo_path') if git_result.artifacts else git_result.summary
             self.logger.info(
-                f"Repository cloned successfully to: {git_result.repo_path}"
+                f"Repository cloned successfully to: {repo_path}"
             )
             return {
-                "repo_path": git_result.repo_path,
+                "repo_path": repo_path,
                 "final_result": "route_to_stack_detect",  # Continue workflow
                 "operation_type": "clone_success",
                 "error_message": None,
@@ -679,12 +695,12 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                 issues_opened = 0
                 error_message = None
                 final_result = f"Unsupported stack detected: {stack}"
-                if issue_result.error_message:
-                    error_message = issue_result.error_message
-                    final_result += f" - Issue creation failed: {issue_result.error_message}"
+                if issue_result.artifacts.get('error_message'):
+                    error_message = issue_result.artifacts.get('error_message')
+                    final_result += f" - Issue creation failed: {issue_result.artifacts.get('error_message')}"
                 else:
                     issues_opened = 1
-                    final_result += f" - Issue created: {issue_result.result}"
+                    final_result += f" - Issue created: {issue_result.summary}"
 
                 return {
                     "stack_detected": stack_detected,
@@ -929,25 +945,71 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                 )
             )
 
-            if issue_result.error_message:
+            issue_error_msg = issue_result.artifacts.get('error_message') if issue_result.artifacts else None
+            if issue_error_msg:
                 self.logger.error(
-                    f"Issue creation failed: {issue_result.error_message}"
+                    f"Issue creation failed: {issue_error_msg}"
                 )
                 return {
-                    "final_result": f"Issue creation failed: {issue_result.error_message}",
+                    "final_result": f"Issue creation failed: {issue_error_msg}",
                     "issues_opened": 0,
                     "operation_type": "issue_error",
                 }
 
             
             if existing_id is None:
-                new_id = self._parse_issue_number(issue_result.result)
+                new_id = self._parse_issue_number(issue_result.summary)
                 if new_id is not None:
                     self._record_issue_id(repo_url, error_type, new_id)
+                    
+                    # Update the run registry with the new issue ID
+                    try:
+                        # Find the current run by repo URL and update with issue ID
+                        current_runs = self.run_registry.find_by_commit_and_repo(repo_url, "manual")
+                        if current_runs:
+                            latest_run = current_runs[0]  # Get most recent
+                            updated = self.run_registry.update(latest_run.run_key, {
+                                'umbrella_issue_id': new_id,
+                                'status': 'FAILED'  # Issue created means workflow failed
+                            })
+                            if updated:
+                                self.logger.info(f"Updated run {latest_run.run_key} with issue ID {new_id}")
+                    except Exception as e:
+                        self.logger.warning(f"Failed to update run registry with issue ID: {e}")
+                    
+                    # If this is a PR merge workflow, update the previous issue with the new issue ID
+                    try:
+                        if self.pr_merge_context:
+                            # Find the previous umbrella issue
+                            previous_issue_id = self.run_registry.find_previous_umbrella_issue(
+                                repo_url, exclude_sha=self.pr_merge_context.get('merged_sha', '')
+                            )
+                            if previous_issue_id:
+                                # Create updated comment linking to the new issue
+                                updated_comment = self._create_issue_link_comment(
+                                    previous_issue_id, 
+                                    self.pr_merge_context.get('merged_sha', 'unknown')[:7],
+                                    new_issue_id=new_id
+                                )
+                                # Comment on the previous issue with the updated link
+                                link_result = self.git_agent.run(
+                                    GitAgentInput(
+                                        query=f"comment on issue {previous_issue_id}: {updated_comment}",
+                                        issue_id=previous_issue_id,
+                                        thread_id=state.get("thread_id")
+                                    )
+                                )
+                                link_error_msg = link_result.artifacts.get('error_message') if link_result.artifacts else None
+                                if link_error_msg:
+                                    self.logger.warning(f"Failed to update previous issue #{previous_issue_id} with new issue link: {link_error_msg}")
+                                else:
+                                    self.logger.info(f"Successfully updated previous issue #{previous_issue_id} with link to new issue #{new_id}")
+                    except Exception as e:
+                        self.logger.warning(f"Error updating previous issue with new issue link: {e}")
 
             self.logger.info("GitHub issue created successfully")
             return {
-                "final_result": f"R2D workflow failed, issue created: {issue_result.result}",
+                "final_result": f"R2D workflow failed, issue created: {issue_result.summary}",
                 "issues_opened": 1,
                 "operation_type": "issue_success",
                 "error_message": None,
@@ -1095,22 +1157,22 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                 )
             )
 
-            if issue_result.error_message:
-                self.logger.error(f"Issue creation failed: {issue_result.error_message}")
+            if issue_result.artifacts.get('error_message'):
+                self.logger.error(f"Issue creation failed: {issue_result.artifacts.get('error_message')}")
                 return {
-                    "final_result": f"Issue creation failed: {issue_result.error_message}",
+                    "final_result": f"Issue creation failed: {issue_result.artifacts.get('error_message')}",
                     "issues_opened": 0,
                     "operation_type": "issue_error",
                 }
 
             # Track new issue ID for deduplication
             if existing_id is None:
-                new_id = self._parse_issue_number(issue_result.result)
+                new_id = self._parse_issue_number(issue_result.summary)
                 if new_id is not None:
                     self._record_issue_id(repo_url, error_type, new_id)
 
             self.logger.info("GitHub issue created successfully")
-            print(f"\n✅ Success! GitHub issue created: {issue_result.result}")
+            print(f"\n✅ Success! GitHub issue created: {issue_result.summary}")
             
             return {
                 "final_result": f"R2D workflow failed, issue created: {issue_result.result}",
@@ -1370,6 +1432,137 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                 "error_message": str(e),
             }
 
+    # --- PR Merge Handling Methods ---
+
+    def _generate_dynamic_branch_name(self) -> str:
+        """Generate timestamp-based branch name."""
+        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+        return f"r2d-{timestamp}"
+
+    def _detect_pr_merge_context(self) -> Optional[Dict[str, Any]]:
+        """
+        Detect PR merge context from environment variables.
+        
+        Returns:
+            PR merge context dictionary if detected, None otherwise
+        """
+        try:
+            pr_context_env = os.environ.get('PR_MERGE_CONTEXT')
+            if pr_context_env:
+                import json
+                pr_context = json.loads(pr_context_env)
+                self.logger.info(f"PR merge context detected: {pr_context}")
+                return pr_context
+        except Exception as e:
+            self.logger.warning(f"Failed to parse PR merge context: {e}")
+        
+        return None
+
+    def _handle_pr_merge_workflow(self, repo_url: str, commit_sha: str) -> Optional[Dict[str, Any]]:
+        """
+        Handle PR merge workflow by detecting previous umbrella issues and linking them.
+        
+        Args:
+            repo_url: Repository URL
+            commit_sha: New commit SHA from PR merge
+            
+        Returns:
+            Dictionary with previous issue information if found, None otherwise
+        """
+        try:
+            # Find previous umbrella issue for this repository (excluding current SHA)
+            previous_issue_id = self.run_registry.find_previous_umbrella_issue(
+                repo_url, exclude_sha=commit_sha
+            )
+            
+            if previous_issue_id:
+                self.logger.info(f"Found previous umbrella issue #{previous_issue_id} for {repo_url}")
+                
+                # Create run for new commit and link to previous issue
+                job_name = f"pr-merge-{self.pr_merge_context.get('pr_number', 'unknown')}"
+                new_run_key = self.run_registry.create_run(
+                    repo_url=repo_url,
+                    commit_sha=commit_sha,
+                    job_name=job_name,
+                    thread_id=str(self.pr_merge_context.get('pr_number', 'unknown'))
+                )
+                
+                # Find the previous run to link them
+                previous_run = self.run_registry.find_latest_run_with_issue(repo_url)
+                if previous_run:
+                    self.run_registry.link_predecessor_run(new_run_key, previous_run.run_key)
+                    self.run_registry.close_old_umbrella_issue(
+                        previous_run.run_key, 
+                        new_issue_id=0,  # Will be updated when new issue is created
+                        new_commit_sha=commit_sha
+                    )
+                    
+                    # Comment on the previous issue to link it to the new workflow
+                    try:
+                        link_comment = self._create_issue_link_comment(
+                            previous_issue_id, commit_sha, new_issue_id=None
+                        )
+                        # Use the standard issue format but with issue_id to trigger commenting
+                        comment_result = self.git_agent.run(
+                            GitAgentInput(
+                                query=f"open issue PR Merge Update for repository {repo_url}: {link_comment}",
+                                issue_id=previous_issue_id,
+                                thread_id=str(self.pr_merge_context.get('pr_number', 'unknown'))
+                            )
+                        )
+                        error_message = comment_result.artifacts.get('error_message') if comment_result.artifacts else None
+                        if error_message:
+                            self.logger.warning(f"Failed to comment on previous issue #{previous_issue_id}: {error_message}")
+                        else:
+                            self.logger.info(f"Successfully commented on previous issue #{previous_issue_id}")
+                    except Exception as e:
+                        self.logger.warning(f"Error commenting on previous issue #{previous_issue_id}: {e}")
+                
+                return {
+                    "previous_issue_id": previous_issue_id,
+                    "new_run_key": new_run_key,
+                    "previous_run": previous_run,
+                    "commit_sha": commit_sha[:7]
+                }
+        except Exception as e:
+            self.logger.error(f"Error handling PR merge workflow: {e}")
+        
+        return None
+
+    def _create_issue_link_comment(self, previous_issue_id: int, new_commit_sha: str, new_issue_id: Optional[int] = None) -> str:
+        """
+        Create a comment to link old issue to new issue for PR merge workflow.
+        
+        Args:
+            previous_issue_id: ID of the previous umbrella issue
+            new_commit_sha: New commit SHA that triggered the workflow
+            new_issue_id: New issue ID if created
+            
+        Returns:
+            Comment text for linking issues
+        """
+        short_sha = new_commit_sha[:7]
+        
+        if new_issue_id:
+            comment = f"""🔄 **New commit detected: `{short_sha}`**
+
+A new commit has been merged, opening fresh pipeline in Issue #{new_issue_id}.
+
+**Previous Issue Status:** Resolved/Follow-up - superseded by new commit
+**New Pipeline:** Issue #{new_issue_id} tracks the latest deployment
+
+This issue is now closed and linked forward to the new deployment pipeline."""
+        else:
+            comment = f"""🔄 **New commit detected: `{short_sha}`**
+
+A new commit has been merged, attempting to open fresh pipeline.
+
+**Previous Issue Status:** Resolved/Follow-up - superseded by new commit
+
+This issue is now closed. A new issue will be created for the latest deployment pipeline."""
+        
+        return comment
+
     # --- Utility Methods (preserved from original implementation) ---
 
     def _generate_dynamic_branch_name(self) -> str:
@@ -1721,7 +1914,8 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                         )
                     )
                 )
-                if not issue_result.error_message:
+                startup_error_msg = issue_result.artifacts.get('error_message') if issue_result.artifacts else None
+                if not startup_error_msg:
                     issues_opened = 1
             except Exception as e:
                 self.logger.error(
@@ -1757,6 +1951,55 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
             thread_id=thread_id,
         )
 
+        # Handle PR merge context and old issue linking
+        pr_merge_results = None
+        current_commit_sha = None
+        run_key = None
+        
+        # Extract commit SHA from PR merge context or generate from current state
+        if self.pr_merge_context:
+            current_commit_sha = self.pr_merge_context.get('merged_sha')
+            self.logger.info(f"Using PR merge SHA: {current_commit_sha}")
+            
+            # Handle previous issue linking for PR merges
+            pr_merge_results = self._handle_pr_merge_workflow(repo_url, current_commit_sha)
+            if pr_merge_results:
+                self.logger.info(f"PR merge workflow handled: {pr_merge_results}")
+        else:
+            # Generate a placeholder SHA for non-PR runs (this would be replaced by actual git SHA in real deployment)
+            current_commit_sha = f"manual-{datetime.now().strftime('%Y%m%d%H%M%S')}"
+            
+        # Create registry entry for this run
+        try:
+            job_name = self.pr_merge_context.get('pr_title', 'R2D Workflow') if self.pr_merge_context else 'Manual R2D Workflow'
+            run_key = self.run_registry.create_run(
+                repo_url=repo_url,
+                commit_sha=current_commit_sha,
+                job_name=job_name,
+                thread_id=thread_id,
+                branch_name=branch_name
+            )
+            self.logger.info(f"Created registry entry with run key: {run_key}")
+            
+            # Link to predecessor run if this is a PR merge
+            if pr_merge_results and pr_merge_results.get('previous_run_key'):
+                linked = self.run_registry.link_predecessor_run(
+                    run_key, pr_merge_results['previous_run_key']
+                )
+                if linked:
+                    self.logger.info(f"Linked to predecessor run: {pr_merge_results['previous_run_key']}")
+                    
+        except Exception as e:
+            self.logger.warning(f"Failed to create registry entry: {e}")
+
+        # Handle PR merge workflow context if detected
+        pr_merge_info = None
+        if self.pr_merge_context:
+            commit_sha = self.pr_merge_context.get('merged_sha', 'unknown')
+            pr_merge_info = self._handle_pr_merge_workflow(repo_url, commit_sha)
+            if pr_merge_info:
+                self.logger.info(f"PR merge workflow initiated: {pr_merge_info['commit_sha']} follows {pr_merge_info['previous_issue_id']}")
+
         # Create initial state
         initial_state: SupervisorAgentState = {
             "input_message": HumanMessage(
@@ -1850,7 +2093,7 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                     "network" in error_message.lower()
                     or "connection" in error_message.lower()
                 ):
-                    error_type = "network_error"
+                                       error_type = "network_error"
                 elif "timeout" in error_message.lower():
                     error_type = "timeout_error"
                 elif (
@@ -1894,17 +2137,19 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
                         issue_id=existing_id,
                     )
                     issue_result = self.git_agent.run(git_input)
-                    if not issue_result.error_message:
+                    error_msg_from_artifacts = issue_result.artifacts.get('error_message') if issue_result.artifacts else None
+                    if not error_msg_from_artifacts:
                         issues_opened = 1
                         # Record new issue id if created
                         if existing_id is None:
-                            new_id = self._parse_issue_number(issue_result.result)
+                            new_id = self._parse_issue_number(issue_result.summary)
                             if new_id is not None:
                                 self._record_issue_id(repo_url, error_type, new_id)
 
-                        self.logger.info(f"Successfully created GitHub issue for workflow failure: {issue_result.result}")
+                        self.logger.info(f"Successfully created GitHub issue for workflow failure: {issue_result.summary}")
                     else:
-                        self.logger.warning(f"Failed to create GitHub issue for workflow failure: {issue_result.error_message}")
+                        issue_error_msg = issue_result.artifacts.get('error_message') if issue_result.artifacts else "Unknown error"
+                        self.logger.warning(f"Failed to create GitHub issue for workflow failure: {issue_error_msg}")
                     
 
             except Exception as issue_error:
@@ -1937,6 +2182,90 @@ Important: Only use routing tokens if the input contains actionable R2D workflow
             )
             return output
 
+    def resume_workflow(self, run_key: str, commit_sha: str) -> SupervisorAgentOutput:
+        """
+        Resume a workflow from a previous state using the run registry.
+        
+        Args:
+            run_key: The unique identifier for the workflow run to resume
+            commit_sha: The commit SHA to use for the resumed workflow
+            
+        Returns:
+            SupervisorAgentOutput: Result of the resumed workflow
+        """
+        self.logger.info(f"Resuming workflow for run_key: {run_key}, commit_sha: {commit_sha}")
+        
+        try:
+            # Get the existing run from registry
+            existing_run = self.run_registry.get_run(run_key)
+            if not existing_run:
+                self.logger.error(f"No run found for key: {run_key}")
+                return SupervisorAgentOutput(
+                    repo_url="unknown",
+                    branch_created=False,
+                    branch_name="unknown",
+                    stack_detected={},
+                    terraform_summary=None,
+                    unsupported=False,
+                    issues_opened=0,
+                    success=False,
+                    message=f"No existing run found for key: {run_key}"
+                )
+            
+            # Update run status to IN_PROGRESS
+            self.run_registry.update(run_key, {
+                'status': 'IN_PROGRESS',
+                'commit_sha': commit_sha,
+                'wait_reason': None
+            })
+            
+            # Create a new SupervisorAgentInput based on the existing run data
+            resume_input = SupervisorAgentInput(
+                repo_url=existing_run.repo_url,
+                branch_name=existing_run.branch_name,
+                thread_id=existing_run.thread_id,
+                dry_run=False  # Assume production run for resume
+            )
+            
+            self.logger.info(f"Resuming workflow for repo: {existing_run.repo_url}")
+            
+            # Run the workflow with the restored input
+            # The workflow will naturally continue from where it left off
+            # based on the state preserved in memory and registry
+            result = self.run(resume_input)
+            
+            # Update registry with final status
+            final_status = 'COMPLETED' if result.success else 'FAILED'
+            self.run_registry.update(run_key, {
+                'status': final_status,
+                'issues_opened': result.issues_opened,
+                'terraform_summary': result.terraform_summary
+            })
+            
+            self.logger.info(f"Workflow resumption {'successful' if result.success else 'failed'} for run_key: {run_key}")
+            return result
+            
+        except Exception as e:
+            self.logger.error(f"Error resuming workflow for run_key {run_key}: {e}")
+            
+            # Update registry with error status
+            self.run_registry.update(run_key, {
+                'status': 'FAILED',
+                'wait_reason': f"Resume error: {str(e)}"
+            })
+            
+            return SupervisorAgentOutput(
+                repo_url=existing_run.repo_url if 'existing_run' in locals() else "unknown",
+                branch_created=False,
+                branch_name=existing_run.branch_name if 'existing_run' in locals() else "unknown",
+                stack_detected={},
+                terraform_summary=None,
+                unsupported=False,
+                issues_opened=0,
+                success=False,
+                message=f"Workflow resumption failed: {str(e)}"
+            )
+
     # --- Memory and Conversation Management ---
 
     def get_conversation_history(self) -> List[Dict[str, any]]: