diff --git a/agent/robot/manager/manager.go b/agent/robot/manager/manager.go index f036e538..cbc10631 100644 --- a/agent/robot/manager/manager.go +++ b/agent/robot/manager/manager.go @@ -138,6 +138,16 @@ func (m *Manager) Start() error { return fmt.Errorf("failed to load robots: %w", err) } + // Set completion callback to clean up ExecutionController when execution finishes + m.pool.SetOnComplete(func(execID, memberID string, status types.ExecStatus) { + // Remove from ExecutionController (cleans up in-memory tracking) + m.execController.Untrack(execID) + // Remove from robot's in-memory execution list + if robot := m.cache.Get(memberID); robot != nil { + robot.RemoveExecution(execID) + } + }) + // Start worker pool if err := m.pool.Start(); err != nil { return fmt.Errorf("failed to start pool: %w", err) diff --git a/agent/robot/pool/pool.go b/agent/robot/pool/pool.go index 9b06585a..70530898 100644 --- a/agent/robot/pool/pool.go +++ b/agent/robot/pool/pool.go @@ -32,18 +32,23 @@ func DefaultConfig() *Config { // ExecutorFactory creates an executor based on the mode type ExecutorFactory func(mode types.ExecutorMode) types.Executor +// OnCompleteCallback is called when an execution completes (success or failure) +// Parameters: execID, memberID, status +type OnCompleteCallback func(execID, memberID string, status types.ExecStatus) + // Pool implements types.Pool interface // Manages a pool of workers that execute robot jobs from a priority queue type Pool struct { - size int // number of workers - queue *PriorityQueue // priority queue for pending jobs - executor types.Executor // default executor for running jobs - executorFactory ExecutorFactory // optional: factory for mode-specific executors - workers []*Worker // worker goroutines - running atomic.Int32 // number of currently running jobs - wg sync.WaitGroup // wait group for graceful shutdown - started bool // whether pool has been started - mu sync.RWMutex // protects started flag + size int // number of workers + queue *PriorityQueue // priority queue for pending jobs + executor types.Executor // default executor for running jobs + executorFactory ExecutorFactory // optional: factory for mode-specific executors + onComplete OnCompleteCallback // optional: callback when execution completes + workers []*Worker // worker goroutines + running atomic.Int32 // number of currently running jobs + wg sync.WaitGroup // wait group for graceful shutdown + started bool // whether pool has been started + mu sync.RWMutex // protects started flag } // New creates a new pool instance with default configuration @@ -86,6 +91,12 @@ func (p *Pool) SetExecutorFactory(factory ExecutorFactory) { p.executorFactory = factory } +// SetOnComplete sets the callback for execution completion +// Called when an execution finishes (completed, failed, or cancelled) +func (p *Pool) SetOnComplete(callback OnCompleteCallback) { + p.onComplete = callback +} + // GetExecutor returns the appropriate executor for the given mode // If factory is set and mode is specified, uses factory; otherwise uses default func (p *Pool) GetExecutor(mode types.ExecutorMode) types.Executor { diff --git a/agent/robot/pool/worker.go b/agent/robot/pool/worker.go index b9e25c38..6889521a 100644 --- a/agent/robot/pool/worker.go +++ b/agent/robot/pool/worker.go @@ -92,12 +92,25 @@ func (w *Worker) execute(item *QueueItem) { } fmt.Printf("Worker %d: Execution failed for robot %s: %v\n", w.id, item.Robot.MemberID, err) + // Notify completion callback with appropriate status + if w.pool.onComplete != nil { + // Determine status based on error type + status := types.ExecFailed + if err == types.ErrExecutionCancelled { + status = types.ExecCancelled + } + w.pool.onComplete(item.ExecID, item.Robot.MemberID, status) + } return } if execution != nil { fmt.Printf("Worker %d: Execution %s completed for robot %s (status: %s)\n", w.id, execution.ID, item.Robot.MemberID, execution.Status) + // Notify completion callback + if w.pool.onComplete != nil { + w.pool.onComplete(execution.ID, item.Robot.MemberID, execution.Status) + } } }