diff --git a/src/cai/cli.py b/src/cai/cli.py index b5e678d6..e1733af4 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -444,6 +444,10 @@ def run_cai_cli( # Reset simple agent manager for clean start from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER AGENT_MANAGER.reset_registry() + + # Register the starting agent with AGENT_MANAGER + starting_agent_name = getattr(starting_agent, "name", last_agent_type) + AGENT_MANAGER.switch_to_single_agent(starting_agent, starting_agent_name) # Initialize command completer and key bindings command_completer = FuzzyCommandCompleter() @@ -552,9 +556,46 @@ def run_cai_cli( # Update parallel_count to reflect changes from /parallel command parallel_count = int(os.getenv("CAI_PARALLEL", "1")) + if current_agent_type != last_agent_type: + # Check if the /agent command already handled the switch + if os.environ.get("CAI_AGENT_SWITCH_HANDLED") == "1": + os.environ["CAI_AGENT_SWITCH_HANDLED"] = "0" # Reset flag + + # Just get the existing agent that was already switched + from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER + + # First try to get the strong reference if available + if hasattr(AGENT_MANAGER, '_current_agent_strong_ref'): + agent = AGENT_MANAGER._current_agent_strong_ref + # Clear the strong reference after using it + delattr(AGENT_MANAGER, '_current_agent_strong_ref') + else: + agent = AGENT_MANAGER.get_active_agent() + + if agent: + last_agent_type = current_agent_type + else: + # If the agent is None (weak reference expired), recreate it + agent = get_agent_by_name(current_agent_type, agent_id="P1") + last_agent_type = current_agent_type + # Re-register with AGENT_MANAGER + agent_name = agent.name if hasattr(agent, "name") else current_agent_type + AGENT_MANAGER.set_active_agent(agent, agent_name, "P1") + continue + try: - # Import is already at the top level + # CRITICAL: Set up history transfer BEFORE creating the new agent + from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER + + # Get the current agent's history before switching + if hasattr(agent, "name"): + current_agent_name = agent.name + current_history = AGENT_MANAGER.get_message_history(current_agent_name) + if current_history: + AGENT_MANAGER._pending_history_transfer = list(current_history) # Make a copy + + # Now create the new agent agent = get_agent_by_name(current_agent_type, agent_id="P1") last_agent_type = current_agent_type @@ -563,11 +604,20 @@ def run_cai_cli( COST_TRACKER.reset_agent_costs() # Use the new switch_to_single_agent method for proper cleanup - from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER - agent_name = getattr(agent, "name", current_agent_type) + # IMPORTANT: Always use the agent's proper name, not the agent key + agent_name = agent.name if hasattr(agent, "name") else current_agent_type - # Switch to the new agent - AGENT_MANAGER.switch_to_single_agent(agent, agent_name) + # Check if this agent has already been switched to by /agent command + # If so, don't switch again to avoid clearing the history + current_active_agent = AGENT_MANAGER.get_active_agent() + current_active_name = AGENT_MANAGER._active_agent_name + + if current_active_name == agent_name: + # Just update the agent reference + agent = current_active_agent + else: + # Switch to the new agent + AGENT_MANAGER.switch_to_single_agent(agent, agent_name) # Sync the model's history with AGENT_MANAGER's history # This ensures the model has its own history from AGENT_MANAGER @@ -1702,7 +1752,8 @@ def main(): # Use the switch_to_single_agent method for proper initialization from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER - agent_name = getattr(agent, "name", agent_type) + # IMPORTANT: Always use the agent's proper name, not the agent key + agent_name = agent.name if hasattr(agent, "name") else agent_type AGENT_MANAGER.switch_to_single_agent(agent, agent_name) # Configure model flags to work well with CLI diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index 372c8576..4c0fda81 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -517,14 +517,13 @@ class AgentCommand(Command): # selected_agent_key was already set above in the agent selection logic pass - # Set the agent key in environment variable (not the agent name) - # Note: selected_agent_key should be defined by now either from regular agent selection - # or from swarm pattern handling - # IMPORTANT: Don't set CAI_AGENT_TYPE for parallel patterns as they don't change the current agent + # IMPORTANT: Don't set CAI_AGENT_TYPE yet - we need to set up history transfer first + # Store the selected agent key for later use + agent_type_to_set = None if 'selected_agent_key' in locals() and not (hasattr(agent, "_pattern") and hasattr(agent._pattern, "type") and str(getattr(agent._pattern.type, 'value', agent._pattern.type)) == "parallel"): - os.environ["CAI_AGENT_TYPE"] = selected_agent_key + agent_type_to_set = selected_agent_key # IMPORTANT: Ensure agent_name is correctly set for the selected agent # This fixes the issue where swarm pattern's agent name lingers @@ -583,6 +582,7 @@ class AgentCommand(Command): break else: # We're switching from single agent to another single agent (or from swarm pattern) + # First check if there's a pending history transfer if hasattr(AGENT_MANAGER, '_pending_history_transfer') and AGENT_MANAGER._pending_history_transfer: current_history = AGENT_MANAGER._pending_history_transfer @@ -646,7 +646,20 @@ class AgentCommand(Command): new_agent = get_agent_by_name(selected_agent_key, agent_id="P1") new_agent_name = getattr(new_agent, "name", selected_agent_key) AGENT_MANAGER.switch_to_single_agent(new_agent, new_agent_name) + + # IMPORTANT: Store a strong reference to prevent garbage collection + # The CLI will pick this up and use it + AGENT_MANAGER._current_agent_strong_ref = new_agent + + # Check if history was transferred + transferred_history = AGENT_MANAGER.get_message_history(new_agent_name) + # NOW set the environment variable AFTER history transfer is complete + if agent_type_to_set: + os.environ["CAI_AGENT_TYPE"] = agent_type_to_set + # Set a flag to tell CLI not to switch again + os.environ["CAI_AGENT_SWITCH_HANDLED"] = "1" + # Double-check agent_name is correct before displaying # This ensures we show the correct agent name even after switching from patterns final_agent_name = agent_name diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index e7ae54e6..82af73f4 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -416,19 +416,16 @@ class OpenAIChatCompletionsModel(Model): if self.agent_name not in AGENT_MANAGER._message_history: AGENT_MANAGER._message_history[self.agent_name] = self.message_history - # Register with SimpleAgentManager only when explicitly created - # This prevents phantom instances during module imports - if agent_id is not None: # Only register when explicitly given an ID - # For parallel agents, register as parallel - if agent_id.startswith("P") and int(os.getenv("CAI_PARALLEL", "1")) > 1: - AGENT_MANAGER.set_parallel_agent(agent_id, self, self.agent_name) - else: - AGENT_MANAGER.set_active_agent(self, self.agent_name, agent_id) - - # CRITICAL: Ensure AGENT_MANAGER uses the same list reference as the model - # This is necessary for proper history clearing to work - if not PARALLEL_ISOLATION.is_parallel_mode(): - AGENT_MANAGER._message_history[self.agent_name] = self.message_history + # NOTE: Models should NOT register themselves with AGENT_MANAGER + # The agent that owns this model will handle registration + # This prevents duplicate registrations with agent keys + + # CRITICAL: Ensure AGENT_MANAGER uses the same list reference as the model + # This is necessary for proper history clearing to work + if agent_id is not None and not PARALLEL_ISOLATION.is_parallel_mode(): + if self.agent_name in AGENT_MANAGER._message_history: + # Share the same list reference + self.message_history = AGENT_MANAGER._message_history[self.agent_name] # Instance-based converter self._converter = _Converter() diff --git a/src/cai/sdk/agents/simple_agent_manager.py b/src/cai/sdk/agents/simple_agent_manager.py index 0ea4a5f4..11086644 100644 --- a/src/cai/sdk/agents/simple_agent_manager.py +++ b/src/cai/sdk/agents/simple_agent_manager.py @@ -25,6 +25,11 @@ class SimpleAgentManager: def set_active_agent(self, agent, agent_name: str, agent_id: str = None): """Set the active agent instance.""" + # CRITICAL: Always use the agent's proper name, not the agent key + # This prevents duplicate registrations like "blueteam_agent" and "Blue Team Agent" + if hasattr(agent, 'name') and agent.name: + agent_name = agent.name + # In single agent mode, use switch_to_single_agent for proper cleanup if not self._parallel_agents and not agent_id: # If we're in single agent mode and no explicit ID is provided @@ -134,22 +139,12 @@ class SimpleAgentManager: # In single agent mode if not self._parallel_agents: - # Always show the active agent, even if it has no history + # In single agent mode, ONLY show the ONE active agent if self._active_agent_name and self._active_agent_name in self._agent_registry: agent_id = self._agent_registry[self._active_agent_name] history = self._message_history.get(self._active_agent_name, []) result[f"{self._active_agent_name} [{agent_id}]"] = history - - # Show all other registered agents that have history - for agent_name, agent_id in sorted(self._agent_registry.items()): - # Skip the active agent (already added above) - if agent_name == self._active_agent_name: - continue - - history = self._message_history.get(agent_name, []) - # Only include non-active agents if they have history - if history: - result[f"{agent_name} [{agent_id}]"] = history + # That's it - no other agents in single agent mode else: # In parallel mode, show all registered agents for agent_name, agent_id in sorted(self._agent_registry.items()): @@ -187,6 +182,11 @@ class SimpleAgentManager: def set_parallel_agent(self, agent_id: str, agent, agent_name: str): """Register a parallel agent.""" + # CRITICAL: Always use the agent's proper name, not the agent key + # This prevents duplicate registrations like "blueteam_agent" and "Blue Team Agent" + if hasattr(agent, 'name') and agent.name: + agent_name = agent.name + # Check if this ID is already registered to a different agent existing_agent_name = self.get_agent_by_id(agent_id) if existing_agent_name and existing_agent_name != agent_name: @@ -290,24 +290,22 @@ class SimpleAgentManager: if self._parallel_agents: return # Only cleanup in single agent mode - # Find all agents with P1 ID - p1_agents = [(name, aid) for name, aid in list(self._agent_registry.items()) if aid == "P1"] + # In single agent mode, there should be ONLY ONE agent + if not self._active_agent_name: + return - if len(p1_agents) <= 1: - return # No duplicates + # Remove ALL agents except the active one + agents_to_remove = [] + for agent_name in list(self._agent_registry.keys()): + if agent_name != self._active_agent_name: + agents_to_remove.append(agent_name) - # Use the tracked active agent name - active_agent_name = self._active_agent_name + for agent_name in agents_to_remove: + del self._agent_registry[agent_name] + if agent_name in self._message_history: + del self._message_history[agent_name] - # Keep only the active agent and those with message history - for agent_name, agent_id in p1_agents: - if agent_name != active_agent_name: - # Check if this agent has any message history - if not self._message_history.get(agent_name): - # No history, safe to remove - del self._agent_registry[agent_name] - if agent_name in self._message_history: - del self._message_history[agent_name] + return def _cleanup_duplicate_ids(self): """Clean up agents with duplicate IDs in parallel mode.""" @@ -369,6 +367,11 @@ class SimpleAgentManager: def switch_to_single_agent(self, agent, agent_name: str): """Switch to a new single agent, properly cleaning up the previous one.""" + # CRITICAL: Always use the agent's proper name, not the agent key + # This prevents duplicate registrations like "blueteam_agent" and "Blue Team Agent" + if hasattr(agent, 'name') and agent.name: + agent_name = agent.name + # Check for pending history transfer (from parallel mode) # This is ONLY used when switching from parallel to single agent mode transfer_history = None @@ -435,12 +438,15 @@ class SimpleAgentManager: self._message_history[agent_name] = [] else: # Agent already has a history entry - # If there's a transfer_history and the current history is empty, use the transfer - if transfer_history and not self._message_history[agent_name]: + # If there's a transfer_history, always use it (this is an explicit transfer request) + if transfer_history: self._message_history[agent_name] = transfer_history # Reset ID counter for cleanliness self._id_counter = 1 + + # Final cleanup to ensure only one agent in single mode + self._cleanup_single_agent_duplicates() def share_swarm_history(self, agent1_name: str, agent2_name: str): """Share message history between two swarm agents.