diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 120000
index 0000000..ac534a3
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1 @@
+AGENT.md
\ No newline at end of file
diff --git a/GEMINI.md b/GEMINI.md
new file mode 120000
index 0000000..681311e
--- /dev/null
+++ b/GEMINI.md
@@ -0,0 +1 @@
+CLAUDE.md
\ No newline at end of file
diff --git a/src/components.rs b/src/components.rs
index e7616ac..0619ee7 100644
--- a/src/components.rs
+++ b/src/components.rs
@@ -39,10 +39,10 @@ pub struct Captured {
 // New component for the tractor beam visual effect
 #[derive(Component)]
 pub struct TractorBeam {
-    pub target: Entity, // The entity being targeted (usually player)
-    pub timer: Timer,   // How long the beam lasts
-    pub width: f32,     // Visual width of the beam
-    pub active: bool,   // Whether the beam is currently active
+    pub target: Option<Entity>, // The entity being targeted (usually player)
+    pub timer: Timer,           // How long the beam lasts
+    pub width: f32,             // Visual width of the beam
+    pub active: bool,           // Whether the beam is currently active
 }
 
 #[derive(Component)]
@@ -50,6 +50,12 @@ pub struct FormationTarget {
     pub position: Vec3,
 }
 
+// Component to store the original formation position for enemies that need to return
+#[derive(Component)]
+pub struct OriginalFormationPosition {
+    pub position: Vec3,
+}
+
 // Enum defining different ways an enemy can attack
 #[derive(Component, Clone, Copy, PartialEq, Debug)]
 pub enum AttackPattern {
@@ -65,6 +71,7 @@ pub enum EnemyState {
     Entering,                 // Flying onto the screen towards formation target
     InFormation,              // Holding position in the formation
     Attacking(AttackPattern), // Diving towards the player using a specific pattern
+    ReturningWithCaptive,     // Boss returning to formation with captured player
 }
 
 #[derive(Component)]
diff --git a/src/enemy.rs b/src/enemy.rs
index a497e15..8433a7b 100644
--- a/src/enemy.rs
+++ b/src/enemy.rs
@@ -1,7 +1,7 @@
 use bevy::prelude::*;
 use std::time::Duration;
 
-use crate::components::{Enemy, EnemyBullet, EnemyState, EnemyType, FormationTarget, TractorBeam, Captured};
+use crate::components::{Enemy, EnemyBullet, EnemyState, EnemyType, FormationTarget, TractorBeam, Captured, OriginalFormationPosition};
 use crate::constants::{ // Only keeping used constants
     ENEMY_BULLET_SIZE, ENEMY_SIZE, ENEMY_SPEED, WINDOW_HEIGHT, WINDOW_WIDTH,
     TRACTOR_BEAM_WIDTH, TRACTOR_BEAM_DURATION, TRACTOR_BEAM_COLOR, CAPTURE_DURATION,
@@ -52,7 +52,7 @@ pub fn spawn_enemies(
 
         // Determine enemy type - now with a chance to spawn Boss enemies
         // Higher stages have a slightly higher boss chance
-        let boss_chance = 0.05 + (stage.number as f32 * 0.01).min(0.15);
+        let boss_chance = 0.30 + (stage.number as f32 * 0.05).min(0.50); // Increased for testing
         let enemy_type = if fastrand::f32() < boss_chance {
             println!("Spawning a Boss enemy!");
             EnemyType::Boss
@@ -113,7 +113,7 @@ pub fn move_enemies(
         (With<Enemy>, With<FormationTarget>),
     >,
     mut attacking_query: Query<
-        (Entity, &mut Transform, &EnemyState, &Enemy), // Add &Enemy here
+        (Entity, &mut Transform, &mut EnemyState, &Enemy, Option<&OriginalFormationPosition>), // Add mutable state and original position
         (With<Enemy>, Without<FormationTarget>),
     >, // Query potential attackers
     time: Res<Time>,
@@ -147,6 +147,10 @@ pub fn move_enemies(
                 transform.translation = target_pos;
                 commands.entity(entity).remove::<FormationTarget>(); // Remove target component
                 *state = EnemyState::InFormation; // Change state
+                
+                // Store the original formation position for potential return
+                commands.entity(entity).insert(OriginalFormationPosition { position: target_pos });
+                
                 println!(
                     "Enemy {:?} reached formation target and is now InFormation.",
                     entity
@@ -159,10 +163,11 @@ pub fn move_enemies(
         }
     }
 
-    // --- Handle Attacking Enemies ---
-    for (entity, mut transform, state, enemy) in attacking_query.iter_mut() {
+    // --- Handle Attacking and Returning Enemies ---
+    for (entity, mut transform, mut state, enemy, original_pos) in attacking_query.iter_mut() {
         // Check what state the enemy is in
-        if let EnemyState::Attacking(attack_pattern) = state {
+        match state.as_ref() {
+            EnemyState::Attacking(attack_pattern) => {
             // Apply different movement based on enemy type
             match enemy.enemy_type {
                 EnemyType::Grunt => {
@@ -230,7 +235,7 @@ pub fn move_enemies(
                                 if has_beam_query.get(entity).is_err() {
                                     // Spawn tractor beam component on this boss
                                     commands.entity(entity).insert(TractorBeam {
-                                        target: Entity::PLACEHOLDER, // Will be filled in by the boss_capture_attack
+                                        target: None, // Will be filled in by the boss_capture_attack
                                         timer: Timer::new(Duration::from_secs_f32(TRACTOR_BEAM_DURATION), TimerMode::Once),
                                         width: TRACTOR_BEAM_WIDTH,
                                         active: false,
@@ -274,6 +279,27 @@ pub fn move_enemies(
                     }
                 }
             }
+            }
+            EnemyState::ReturningWithCaptive => {
+                // Boss returning to formation with captured player
+                if let Some(original_pos) = original_pos {
+                    let direction = original_pos.position - transform.translation;
+                    let distance = direction.length();
+                    let return_speed = ENEMY_SPEED * speed_multiplier * 0.7; // Slightly slower when carrying captive
+                    
+                    if distance < 5.0 { // Close enough to formation position
+                        // Return to formation
+                        transform.translation = original_pos.position;
+                        *state = EnemyState::InFormation;
+                        println!("Boss {:?} returned to formation with captive!", entity);
+                    } else {
+                        // Move towards formation position
+                        let move_delta = direction.normalize() * return_speed * time.delta_seconds();
+                        transform.translation += move_delta;
+                    }
+                }
+            }
+            _ => {} // Handle other states if needed
         }
 
         // Despawn if off screen (This should be inside the loop)
@@ -371,11 +397,13 @@ pub fn trigger_attack_dives(
             let mut selected_pattern = match enemy_type {
                 // For Boss enemies, occasionally use the CaptureBeam pattern
                 EnemyType::Boss => {
-                    if fastrand::f32() < 0.4 { // 40% chance for Boss to use CaptureBeam
+                    if fastrand::f32() < 0.8 { // 80% chance for Boss to use CaptureBeam (increased for testing)
+                        println!("Boss {:?} selected CaptureBeam attack!", chosen_entity);
                         AttackPattern::CaptureBeam
                     } else {
                         // Otherwise use a random pattern from the stage config
                         let pattern_index = fastrand::usize(..current_config.attack_patterns.len());
+                        println!("Boss {:?} selected {:?} attack!", chosen_entity, current_config.attack_patterns[pattern_index]);
                         current_config.attack_patterns[pattern_index]
                     }
                 },
@@ -456,16 +484,16 @@ pub fn boss_capture_attack(
     time: Res<Time>,
     mut boss_query: Query<(Entity, &Transform, &mut TractorBeam)>,
     player_query: Query<(Entity, &Transform), (With<Player>, Without<Captured>)>, // Only target non-captured players
-    has_beam_query: Query<&TractorBeam>,
+    enemy_query: Query<&EnemyState, With<Enemy>>,
 ) {
     for (boss_entity, boss_transform, mut tractor_beam) in boss_query.iter_mut() {
         // Tick the beam timer
         tractor_beam.timer.tick(time.delta());
         
         // If player exists and beam is not active yet, set player as target
-        if !tractor_beam.active && tractor_beam.target == Entity::PLACEHOLDER {
+        if !tractor_beam.active && tractor_beam.target.is_none() {
             if let Ok((player_entity, _)) = player_query.get_single() {
-                tractor_beam.target = player_entity;
+                tractor_beam.target = Some(player_entity);
                 tractor_beam.active = true;
                 println!("Boss {:?} activated tractor beam targeting player!", boss_entity);
                 
@@ -504,9 +532,15 @@ pub fn boss_capture_attack(
                     // Boss returns to formation with captured player
                     commands.entity(boss_entity).remove::<TractorBeam>();
                     
-                    // TODO: Implement logic for boss to return to formation
-                    // For now, just despawn the boss to simplify
-                    commands.entity(boss_entity).despawn_recursive();
+                    // Change boss state to returning with captive
+                    if let Ok(enemy) = enemy_query.get(boss_entity) {
+                        if let EnemyState::Attacking(_) = enemy {
+                            commands.entity(boss_entity).insert(EnemyState::ReturningWithCaptive);
+                            
+                            // Clean up the beam visual
+                            commands.entity(boss_entity).despawn_descendants();
+                        }
+                    }
                     break;
                 }
             }
@@ -517,12 +551,15 @@ pub fn boss_capture_attack(
             println!("Boss {:?} tractor beam expired", boss_entity);
             commands.entity(boss_entity).remove::<TractorBeam>();
             
-            // If we had a proper visual beam, we'd despawn it here
-            // For now, just make sure we clean up by despawning all children
+            // Clean up the beam visual
             commands.entity(boss_entity).despawn_descendants();
             
-            // For simplicity, after beam attack the boss flies off-screen
-            // In a more complete implementation, it might return to formation
+            // Boss returns to formation after failed capture
+            if let Ok(enemy_state) = enemy_query.get(boss_entity) {
+                if let EnemyState::Attacking(_) = enemy_state {
+                    commands.entity(boss_entity).insert(EnemyState::ReturningWithCaptive);
+                }
+            }
         }
     }
 }
\ No newline at end of file
diff --git a/src/player.rs b/src/player.rs
index 9e7b359..2cc49e1 100644
--- a/src/player.rs
+++ b/src/player.rs
@@ -1,4 +1,5 @@
 use bevy::prelude::*;
+use bevy::ecs::system::ParamSet;
 use std::time::Duration;
 
 use crate::components::{Bullet, Captured, Enemy, Invincible, Player};
@@ -70,51 +71,40 @@ pub fn move_player(
 pub fn handle_captured_player(
     mut commands: Commands,
     time: Res<Time>,
-    player_query: Query<(Entity, &Transform, &Captured)>,
-    mut player_mut_query: Query<(&mut Transform, &mut Captured)>,
-    enemy_query: Query<&Transform, With<Enemy>>,
+    mut set: ParamSet<(
+        Query<(Entity, &mut Transform, &mut Captured)>,
+        Query<&Transform, (With<Enemy>, Without<Player>)>,
+    )>,
     mut lives: ResMut<PlayerLives>,
     mut respawn_timer: ResMut<PlayerRespawnTimer>,
     mut next_state: ResMut<NextState<AppState>>,
 ) {
-    // First, collect data from all captured players
-    let mut to_process = Vec::new();
-
-    for (entity, transform, captured) in player_query.iter() {
-        // Check if the boss exists
-        let boss_exists = enemy_query.get(captured.boss_entity).is_ok();
-        let boss_pos = if boss_exists {
-            enemy_query
-                .get(captured.boss_entity)
-                .map(|t| t.translation)
-                .ok()
-        } else {
-            None
-        };
-
-        // Create a copy of the timer to check if it would finish
-        let mut timer_copy = captured.timer.clone();
-        timer_copy.tick(time.delta());
-
-        to_process.push((
-            entity,
-            transform.translation,
-            boss_pos,
-            timer_copy.finished(),
-        ));
+    // First, collect data about captured players and their bosses
+    let mut captured_data = Vec::new();
+    
+    // Get player data
+    for (entity, transform, captured) in set.p0().iter() {
+        captured_data.push((entity, transform.translation, captured.boss_entity, captured.timer.clone()));
     }
-
-    // Now process each player separately
-    for (entity, current_pos, boss_pos_opt, timer_would_finish) in to_process {
-        if let Ok((mut transform, mut captured)) = player_mut_query.get_mut(entity) {
-            // Tick the real timer
+    
+    // Process each captured player
+    for (player_entity, _player_pos, boss_entity, mut timer) in captured_data {
+        // Check if the boss exists and get its position
+        let boss_pos = set.p1().get(boss_entity).map(|t| t.translation).ok();
+        
+        // Tick the timer
+        timer.tick(time.delta());
+        
+        // Update the player
+        if let Ok((entity, mut transform, mut captured)) = set.p0().get_mut(player_entity) {
+            // Update the actual timer
             captured.timer.tick(time.delta());
-
-            match boss_pos_opt {
+            
+            match boss_pos {
                 Some(boss_pos) => {
                     // Boss exists, update player position
                     let target_pos = boss_pos - Vec3::new(0.0, PLAYER_SIZE.y + 10.0, 0.0);
-                    transform.translation = current_pos.lerp(target_pos, 0.2);
+                    transform.translation = transform.translation.lerp(target_pos, 0.2);
                 }
                 None => {
                     // Boss is gone, release player but lose a life
@@ -127,11 +117,12 @@ pub fn handle_captured_player(
                         &mut next_state,
                         entity,
                     );
+                    continue; // Skip the rest of processing for this player
                 }
             }
-
+            
             // If capture duration expires, player escapes but loses a life
-            if timer_would_finish || captured.timer.finished() {
+            if captured.timer.finished() {
                 println!("Player escaped from capture after timer expired!");
                 commands.entity(entity).remove::<Captured>();
                 lose_life_and_respawn(