diff --git a/src/agents/base_agent.py b/src/agents/base_agent.py index e57d5ab..2ab41dd 100644 --- a/src/agents/base_agent.py +++ b/src/agents/base_agent.py @@ -48,8 +48,8 @@ def clone_repository(self, repo_url: str, target_dir: str) -> bool: return clone_result.exit_code == 0 - def create_agent_conversation(self) -> tuple[Agent, Conversation]: - """Create an agent and conversation instance.""" + def create_agent_conversation(self, secrets: Optional[Dict[str, str]] = None) -> tuple[Agent, Conversation]: + """Create an agent and conversation instance with optional secrets.""" agent = Agent( llm=self.llm, tools=get_default_tools(enable_browser=False), @@ -61,13 +61,15 @@ def create_agent_conversation(self) -> tuple[Agent, Conversation]: visualizer=None, ) + # Add secrets to the conversation if provided + if secrets: + conversation.update_secrets(secrets) + return agent, conversation - def setup_git_environment(self): - """Set up git environment variables.""" - github_token = os.getenv('GITHUB_TOKEN') - if github_token: - self.workspace.execute_command(f"export GITHUB_TOKEN={github_token}") + def get_github_token(self) -> Optional[str]: + """Get GitHub token from environment if available.""" + return os.getenv('GITHUB_TOKEN') def cleanup_directory(self, directory: str): """Clean up a workspace directory.""" diff --git a/src/agents/cve_scanner.py b/src/agents/cve_scanner.py index c6d676e..43a1680 100644 --- a/src/agents/cve_scanner.py +++ b/src/agents/cve_scanner.py @@ -47,7 +47,14 @@ def scan_repository(self, repo_url: str) -> List[Dict[str, Any]]: # Create agent with default tools print("🔍 CVE Scanner Status: Initializing AI agent") - agent, conversation = self.create_agent_conversation() + + # Prepare secrets for the conversation + secrets = {} + github_token = self.get_github_token() + if github_token: + secrets["GITHUB_TOKEN"] = github_token + + agent, conversation = self.create_agent_conversation(secrets if secrets else None) # Send scanning instructions to the agent scan_message = self.prompt_manager.get_scanner_main_prompt( diff --git a/src/agents/cve_solver.py b/src/agents/cve_solver.py index 50c8430..81f1a1c 100644 --- a/src/agents/cve_solver.py +++ b/src/agents/cve_solver.py @@ -62,10 +62,13 @@ def solve_vulnerability(self, vulnerability: Dict[str, Any], repo_url: str, port if activity_callback: activity_callback("Initializing AI agent") - agent, conversation = self.create_agent_conversation() + # Prepare secrets for the conversation + secrets = {} + github_token = self.get_github_token() + if github_token: + secrets["GITHUB_TOKEN"] = github_token - # Set up environment variables for the agent - self.setup_git_environment() + agent, conversation = self.create_agent_conversation(secrets if secrets else None) # Create fix instructions fix_message = self.prompt_manager.get_solver_main_prompt(