fix(swarm): close lifecycle admission on shutdown

internal/swarm/lifecycle.go

@@ -12,6 +12,12 @@ import (
 // team is at its slot cap the member is queued and launches when a slot frees
 // (it stays pending in the coordinator until then).
 func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) {
+	release, err := s.beginLifecycleAdmission()
+	if err != nil {
+		return "", err
+	}
+	defer release()
+
 	def, err := s.registry.Lookup(agentType)
 	if err != nil {
 		return "", err
@@ -23,33 +29,67 @@ func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string
 	}
 	s.rememberCwd(id, cwd)
 	spec := s.buildSpec(pol, id, id, team, def, task, cwd)
-	s.dispatch(spec)
+	s.dispatchAdmitted(spec)
 	return id, nil
 }
 
+// beginLifecycleAdmission admits one unit of lifecycle work if shutdown has not
+// begun. On success it returns a release func the caller MUST call exactly once
+// (a deferred call is the norm) — Close blocks until every admitted unit has
+// released. On shutdown it returns a nil func and ErrSwarmClosed, and there is
+// nothing to release.
+//
+// It deliberately does NOT keep lifecycleMu held for the caller's duration, which
+// an earlier revision did. Close needs the write lock before it can cancel, so a
+// launcher that blocks in Launch until its context is cancelled would deadlock
+// against every Close caller: the launch holds the read lock while Close waits
+// for the write lock to issue the cancellation neither side can reach. Holding the
+// lock only across the flag check and the counter increment keeps the ordering
+// guarantee (no admission once closed is set) without letting external launcher
+// code sit inside the critical section.
+//
+// The read lock is enough to make the WaitGroup safe: Close flips closed under
+// the write lock, so an Add here either completes before Close acquires it (and is
+// therefore visible to the later Wait) or observes closed and is refused.
+func (s *Swarm) beginLifecycleAdmission() (func(), error) {
+	s.lifecycleMu.RLock()
+	defer s.lifecycleMu.RUnlock()
+	if s.closed {
+		return nil, ErrSwarmClosed
+	}
+	s.lifecycleWork.Add(1)
+	return s.lifecycleWork.Done, nil
+}
+
 // dispatch admits a spec to its team (launching now or queuing for a slot).
-func (s *Swarm) dispatch(spec MemberSpec) {
+// The caller must hold an admission ticket from beginLifecycleAdmission.
+func (s *Swarm) dispatchAdmitted(spec MemberSpec) {
 	t := s.team(spec.Team)
 	if t.admit(spec) {
-		s.launch(t, spec)
+		s.launchAdmitted(t, spec)
 	}
 	// Otherwise the spec is queued; the coordinator task stays pending until a
 	// slot frees and afterExit launches it.
 }
 
 // launch starts a member for spec and supervises it. A synchronous launch
 // failure fails the task and frees the slot.
-func (s *Swarm) launch(t *Team, spec MemberSpec) {
+// The caller must hold an admission ticket from beginLifecycleAdmission.
+func (s *Swarm) launchAdmitted(t *Team, spec MemberSpec) {
 	handle, err := s.launcher.Launch(s.baseCtx, spec)
 	if err != nil {
 		_ = s.coord.Fail(spec.TaskID, "launch: "+err.Error())
-		s.afterExit(t)
+		s.afterExitAdmitted(t)
 		return
 	}
 	m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
 	t.addMember(m)
 	_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
-	go s.watch(t, m, spec)
+	s.watchers.Add(1)
+	go func() {
+		defer s.watchers.Done()
+		s.watch(t, m, spec)
+	}()
 }
 
 // watch awaits a member, applies bounded relaunch on temporary failures, records
@@ -60,22 +100,28 @@ func (s *Swarm) launch(t *Team, spec MemberSpec) {
 // life (Member.ID == MemberSpec.ID); a retry never reuses the struct for a
 // different spec.
 func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) {
-	res, err := m.handle.Wait()
-	if err != nil {
-		if isRetryable(err) && m.restarts < maxMemberRestarts {
-			if nh, relErr := s.launcher.Launch(s.baseCtx, spec); relErr == nil {
-				m.restarts++
-				m.handle = nh
-				go s.watch(t, m, spec)
-				return
+	for {
+		res, err := m.handle.Wait()
+		if err != nil {
+			if isRetryable(err) && m.restarts < maxMemberRestarts {
+				if release, admitErr := s.beginLifecycleAdmission(); admitErr == nil {
+					nh, relErr := s.launcher.Launch(s.baseCtx, spec)
+					release()
+					if relErr == nil {
+						m.restarts++
+						m.handle = nh
+						continue
+					}
+				}
+				// Fall through if shutdown started or relaunch failed.
 			}
-			// fall through to record the original error if relaunch fails
+			// res.SessionID is preserved by the handle even on error, so a member that
+			// ran then failed stays drillable; a pure launch error carries an empty id.
+			_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
+		} else {
+			_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
 		}
-		// res.SessionID is preserved by the handle even on error, so a member that
-		// ran then failed stays drillable; a pure launch error carries an empty id.
-		_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
-	} else {
-		_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
+		break
 	}
 	t.removeMember(m.ID)
 	s.afterExit(t)
@@ -85,17 +131,35 @@ func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) {
 // any. Each exit drains at most one queued member; that member's own exit drains
 // the next, so the queue empties one-per-slot without unbounded recursion.
 func (s *Swarm) afterExit(t *Team) {
+	release, err := s.beginLifecycleAdmission()
+	if err != nil {
+		t.releaseSlot()
+		return
+	}
+	defer release()
+	s.afterExitAdmitted(t)
+}
+
+// afterExitAdmitted drains at most one queued spec while lifecycle admission is
+// held open. The caller must hold an admission ticket from beginLifecycleAdmission.
+func (s *Swarm) afterExitAdmitted(t *Team) {
 	next, ok := t.onExit()
 	if !ok {
 		return
 	}
-	s.launch(t, next)
+	s.launchAdmitted(t, next)
 }
 
 // Handoff transfers a task to a fresh member of toAgentType, delivering a note to
 // the new member's inbox and marking the original task handed-off. It returns the
 // new task id. A handoff of an already-terminal task is rejected (fail closed).
 func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) {
+	release, err := s.beginLifecycleAdmission()
+	if err != nil {
+		return "", err
+	}
+	defer release()
+
 	task, ok := s.coord.Get(taskID)
 	if !ok {
 		return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID)
@@ -131,14 +195,20 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string)
 	// Retire the original task (it has been re-delegated).
 	_ = s.coord.SetStatus(taskID, StatusHandedOff)
 	spec := s.buildSpec(pol, newID, newID, team, def, handoffTask, cwd)
-	s.dispatch(spec)
+	s.dispatchAdmitted(spec)
 	return newID, nil
 }
 
 // AdoptOrphans re-parents tasks in a team whose owning member is no longer live
 // (e.g. a crashed worker) onto fresh members of toAgentType, returning the
 // adopted task ids. Terminal tasks and tasks with a live owner are left alone.
 func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) {
+	release, err := s.beginLifecycleAdmission()
+	if err != nil {
+		return nil, err
+	}
+	defer release()
+
 	def, err := s.registry.Lookup(toAgentType)
 	if err != nil {
 		return nil, err
@@ -161,7 +231,7 @@ func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string
 		cwd := s.cwdFor(task.ID)
 		s.rememberCwd(task.ID, cwd)
 		spec := s.buildSpec(pol, newAgent, task.ID, team, def, task.Description, cwd)
-		s.dispatch(spec)
+		s.dispatchAdmitted(spec)
 		adopted = append(adopted, task.ID)
 	}
 	return adopted, nil