Safe Assignment Locking

# Orchestrator loop (simplified)
while not constellation.is_complete():
    ready_tasks = constellation.get_ready_tasks()  # ← May see stale state
    await schedule_ready_tasks(ready_tasks)        # ← Dispatch based on stale view
    await wait_for_task_completion()

async def wait_for_pending_modifications(
    self, timeout: Optional[float] = None
) -> bool:
    """Wait for all pending modifications to complete."""

    if not self._pending_modifications:
        return True

    timeout = timeout or self._modification_timeout
    start_time = asyncio.get_event_loop().time()

    try:
        while self._pending_modifications:
            # Get current pending tasks
            pending_tasks = list(self._pending_modifications.keys())
            pending_futures = list(self._pending_modifications.values())

            self.logger.info(
                f"⏳ Waiting for {len(pending_tasks)} pending modification(s): {pending_tasks}"
            )

            # Calculate remaining timeout
            elapsed = asyncio.get_event_loop().time() - start_time
            remaining_timeout = timeout - elapsed

            if remaining_timeout <= 0:
                raise asyncio.TimeoutError()

            # Wait for all current pending modifications
            await asyncio.wait_for(
                asyncio.gather(*pending_futures, return_exceptions=True),
                timeout=remaining_timeout,
            )

            # Check if new modifications were added during the wait
            if not self._pending_modifications:
                break

            # Small delay to allow new registrations to settle
            await asyncio.sleep(0.01)

        self.logger.info("✅ All pending modifications completed")
        return True

    except asyncio.TimeoutError:
        pending = list(self._pending_modifications.keys())
        self.logger.warning(
            f"⚠️ Timeout waiting for modifications after {timeout}s. "
            f"Proceeding anyway. Pending: {pending}"
        )
        # Clear all pending modifications to prevent permanent deadlock
        self._pending_modifications.clear()
        return False

Algorithm: Safe Assignment Locking and Asynchronous Rescheduling Protocol

Input: Event stream E, current TaskConstellation C
Output: Consistent and updated C with newly scheduled ready tasks

while system is running do
    foreach event e ∈ E do
        if e is TASK_COMPLETED or TASK_FAILED then
            async enqueue(e)  // Record for processing
        end
    end

    acquire(assign_lock)  // Suspend new assignments

    while queue not empty do
        e ← dequeue()
        Δ ← invoke(ConstellationAgent, edit(C, e))  // Propose DAG edits
        C ← apply(C, Δ)                             // Update structure
        validate(C)                                  // Ensure invariants I1-I3
        publish(CONSTELLATION_MODIFIED, t)
        C ← synchronize(C, T_C)  // Merge completed tasks
    end

    release(assign_lock)  // Resume orchestration

    // Rescheduling Phase (outside lock)
    T_R ← get_ready_tasks(C)
    foreach t ∈ T_R do
        async dispatch(t)
        async publish(TASK_STARTED, t)
    end
end

async def _handle_task_event(self, event: TaskEvent) -> None:
    """Handle task completion/failure events."""

    if event.event_type not in [EventType.TASK_COMPLETED, EventType.TASK_FAILED]:
        return

    constellation_id = event.data.get("constellation_id")
    if not constellation_id:
        return

    # Register pending modification
    if event.task_id not in self._pending_modifications:
        modification_future = asyncio.Future()
        self._pending_modifications[event.task_id] = modification_future
        self._stats["total_modifications"] += 1

        self.logger.info(
            f"🔒 Registered pending modification for task '{event.task_id}'"
        )

        # Set timeout to auto-complete if modification takes too long
        asyncio.create_task(
            self._auto_complete_on_timeout(event.task_id, modification_future)
        )

# task_id -> Future mapping
self._pending_modifications: Dict[str, asyncio.Future] = {}

async def _handle_constellation_event(self, event: ConstellationEvent) -> None:
    """Handle constellation modification events."""

    if event.event_type != EventType.CONSTELLATION_MODIFIED:
        return

    task_ids = event.data.get("on_task_id")
    if not task_ids:
        return

    new_constellation = event.data.get("new_constellation")
    if new_constellation:
        self._current_constellation = new_constellation

    # Mark modifications as complete
    for task_id in task_ids:
        if task_id in self._pending_modifications:
            future = self._pending_modifications[task_id]
            if not future.done():
                future.set_result(True)  # Unblocks wait_for_pending_modifications
                self._stats["completed_modifications"] += 1
                self.logger.info(
                    f"✅ Completed modification for task '{task_id}'"
                )
            del self._pending_modifications[task_id]

async def _auto_complete_on_timeout(
    self, task_id: str, future: asyncio.Future
) -> None:
    """Auto-complete a pending modification if it times out."""

    try:
        await asyncio.sleep(self._modification_timeout)  # Default: 600s

        if not future.done():
            self._stats["timeout_modifications"] += 1
            self.logger.warning(
                f"⚠️ Modification for task '{task_id}' timed out. "
                f"Auto-completing to prevent deadlock."
            )
            future.set_result(False)
            if task_id in self._pending_modifications:
                del self._pending_modifications[task_id]
    except asyncio.CancelledError:
        raise

t0: Task A completes → Agent starts editing
t1: Agent modifies constellation (Task A still RUNNING in agent's copy)
t2: Task B completes (orchestrator marks as COMPLETED)
t3: Agent publishes CONSTELLATION_MODIFIED
t4: Orchestrator syncs...

def merge_and_sync_constellation_states(
    self, orchestrator_constellation: TaskConstellation
) -> TaskConstellation:
    """Merge constellation states: structural changes + execution state."""

    if not self._current_constellation:
        return orchestrator_constellation

    # Use agent's constellation as base (has structural modifications)
    merged = self._current_constellation

    # Preserve execution state from orchestrator for existing tasks
    for task_id, orchestrator_task in orchestrator_constellation.tasks.items():
        if task_id in merged.tasks:
            agent_task = merged.tasks[task_id]

            # If orchestrator's state is more advanced, preserve it
            if self._is_state_more_advanced(
                orchestrator_task.status, agent_task.status
            ):
                # Preserve orchestrator's state and results
                agent_task._status = orchestrator_task.status
                agent_task._result = orchestrator_task.result
                agent_task._error = orchestrator_task.error
                agent_task._execution_start_time = orchestrator_task.execution_start_time
                agent_task._execution_end_time = orchestrator_task.execution_end_time

    # Update constellation state
    merged.update_state()

    return merged

def _is_state_more_advanced(self, state1, state2) -> bool:
    """Check if state1 is more advanced than state2."""

    state_levels = {
        TaskStatus.PENDING: 0,
        TaskStatus.WAITING_DEPENDENCY: 1,
        TaskStatus.RUNNING: 2,
        TaskStatus.COMPLETED: 3,
        TaskStatus.FAILED: 3,      # Terminal states equally advanced
        TaskStatus.CANCELLED: 3,
    }

    level1 = state_levels.get(state1, 0)
    level2 = state_levels.get(state2, 0)

    return level1 > level2

async def _sync_constellation_modifications(
    self, constellation: TaskConstellation
) -> TaskConstellation:
    """Synchronize pending constellation modifications."""

    if self._modification_synchronizer:
        # Wait for agent to finish any pending edits
        await self._modification_synchronizer.wait_for_pending_modifications()

        # Merge agent's structural changes with orchestrator's execution state
        constellation = self._modification_synchronizer \
            .merge_and_sync_constellation_states(
                orchestrator_constellation=constellation,
            )

    return constellation

# Process ALL pending modifications in one cycle
while self._pending_modifications:
    # Wait for all to complete
    ...

t0: Task A completes → enqueue(A)
t3: Task B completes → enqueue(B)
t4: Task C completes → enqueue(C)

t5: acquire(lock)
t6: Process A → Δ_A
t7: Process B → Δ_B
t8: Process C → Δ_C
t9: Apply all Δs atomically
t10: release(lock)

from galaxy.constellation.orchestrator import TaskConstellationOrchestrator
from galaxy.session.observers.constellation_sync_observer import (
    ConstellationModificationSynchronizer
)

# Create orchestrator
orchestrator = TaskConstellationOrchestrator(device_manager)

# Create and attach synchronizer
synchronizer = ConstellationModificationSynchronizer(orchestrator)
orchestrator.set_modification_synchronizer(synchronizer)

# Subscribe to events
from galaxy.core.events import get_event_bus
event_bus = get_event_bus()
event_bus.subscribe(synchronizer)

# Orchestrate with synchronization
results = await orchestrator.orchestrate_constellation(constellation)

# Increase timeout for slow LLM responses
synchronizer.set_modification_timeout(1200.0)  # 20 minutes

# Check pending modifications
if synchronizer.has_pending_modifications():
    pending = synchronizer.get_pending_task_ids()
    print(f"Waiting for modifications: {pending}")

# Get statistics
stats = synchronizer.get_statistics()
print(f"Total: {stats['total_modifications']}")
print(f"Completed: {stats['completed_modifications']}")
print(f"Timeouts: {stats['timeout_modifications']}")