fix(daemon): bound pending sessions and prune on completion

internal/daemon/session.go

@@ -32,6 +32,9 @@ const (
 // ErrSessionExists is returned when starting a session ID already in use.
 var ErrSessionExists = errors.New("daemon: session already exists")
 
+// ErrSessionOverloaded means the manager's unfinished-session limit is reached.
+var ErrSessionOverloaded = errors.New("daemon: pending session limit reached")
+
 // ErrSessionNotFound is returned by Get/Attach for an unknown session ID.
 var ErrSessionNotFound = errors.New("daemon: session not found")
 
@@ -219,6 +222,9 @@ func (s *Session) isFinished() bool {
 // daemon does not accumulate finished sessions without limit.
 const defaultMaxSessions = 256
 
+// defaultMaxPending bounds queued and running sessions independently of history.
+const defaultMaxPending = 256
+
 // SessionManager owns the live session registry and routes each session to a
 // pool worker. Mirrors session-manager.js + roster.js. Finished sessions are
 // retained (so a late `attach` sees history) up to MaxSessions, past which the
@@ -227,8 +233,10 @@ type SessionManager struct {
 	pool        *Pool
 	maxBuffer   int
 	maxSessions int
+	maxPending  int
 
 	mu       sync.Mutex
+	pending  int // queued + running; protected by mu
 	sessions map[string]*Session
 	order    []string // creation order, for FIFO eviction of finished sessions
 }
@@ -240,6 +248,9 @@ type SessionManagerOptions struct {
 	// MaxSessions caps the retained session registry; 0 => default. The oldest
 	// FINISHED sessions are evicted once the cap is exceeded.
 	MaxSessions int
+	// MaxPending caps unfinished sessions (queued plus running), independently
+	// of retained history. Non-positive values use the default of 256.
+	MaxPending int
 }
 
 // NewSessionManager builds a manager over pool.
@@ -251,24 +262,39 @@ func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) {
 	if maxSessions <= 0 {
 		maxSessions = defaultMaxSessions
 	}
+	maxPending := opts.MaxPending
+	if maxPending <= 0 {
+		maxPending = defaultMaxPending
+	}
 	return &SessionManager{
 		pool:        opts.Pool,
 		maxBuffer:   opts.MaxBuffer,
 		maxSessions: maxSessions,
+		maxPending:  maxPending,
 		sessions:    map[string]*Session{},
 	}, nil
 }
 
 // Start creates a session for spec and dispatches it to the pool. It returns the
 // Session immediately (non-blocking); the run proceeds in the background and the
 // session transitions queued -> running -> done/failed. A duplicate ID is
-// rejected with ErrSessionExists.
+// rejected with ErrSessionExists; a full admission limit returns
+// ErrSessionOverloaded without creating a session or background goroutine.
 func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, error) {
 	m.mu.Lock()
 	if _, exists := m.sessions[spec.Session]; exists {
 		m.mu.Unlock()
 		return nil, ErrSessionExists
 	}
+	if err := ctx.Err(); err != nil {
+		m.mu.Unlock()
+		return nil, err
+	}
+	if m.pending >= m.maxPending {
+		m.mu.Unlock()
+		return nil, ErrSessionOverloaded
+	}
+	m.pending++
 	sess := newSession(spec.Session, spec.Cwd, m.maxBuffer)
 	m.sessions[spec.Session] = sess
 	m.order = append(m.order, spec.Session)
@@ -277,7 +303,11 @@ func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session,
 
 	go func() {
 		code, err := m.pool.Run(ctx, spec, sess)
+		m.mu.Lock()
 		sess.finish(code, err)
+		m.pending--
+		m.pruneLocked()
+		m.mu.Unlock()
 	}()
 	return sess, nil
 }

internal/daemon/session_test.go

@@ -3,6 +3,8 @@ package daemon
 import (
 	"context"
 	"errors"
+	"fmt"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -208,6 +210,130 @@ func TestSessionManagerKeepsRunningOverCap(t *testing.T) {
 	}
 }
 
+func TestSessionManagerAdmissionBound(t *testing.T) {
+	block := make(chan struct{})
+	launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block})
+	pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
+	mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2})
+	ctx, cancel := context.WithCancel(context.Background())
+	defer func() {
+		cancel()
+		close(block)
+		testutil.WaitFor(t, "all admitted sessions finished", func() bool {
+			mgr.mu.Lock()
+			defer mgr.mu.Unlock()
+			return mgr.pending == 0
+		})
+	}()
+	var wg sync.WaitGroup
+	var accepted atomic.Int32
+	for i := 0; i < 300; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if _, err := mgr.Start(ctx, WorkerSpec{Session: fmt.Sprintf("s%d", i)}); err == nil {
+				accepted.Add(1)
+			}
+		}()
+	}
+	wg.Wait()
+	if got := accepted.Load(); got != 256 {
+		t.Fatalf("accepted %d unfinished sessions, want default bound 256", got)
+	}
+	if got := len(mgr.Statuses()); got != 256 {
+		t.Fatalf("registered %d sessions, want 256", got)
+	}
+}
+
+func TestSessionManagerCancellationReleasesAdmission(t *testing.T) {
+	block := make(chan struct{})
+	var release sync.Once
+	defer release.Do(func() { close(block) })
+	launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block}, &fakeWorker{pid: 2, exitCode: ExitPermanent}, &fakeWorker{pid: 3})
+	pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
+	mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxPending: 2, MaxSessions: 10})
+	first, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	testutil.WaitFor(t, "first running", func() bool { return first.State() == SessionRunning })
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	queued, err := mgr.Start(ctx, WorkerSpec{Session: "queued"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"}); s != nil || !errors.Is(err, ErrSessionOverloaded) {
+		t.Fatalf("overloaded Start = (%v, %v)", s, err)
+	}
+	if _, ok := mgr.Get("retry"); ok {
+		t.Fatal("rejected session was registered")
+	}
+	if _, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"}); !errors.Is(err, ErrSessionExists) {
+		t.Fatalf("duplicate at capacity: %v", err)
+	}
+	cancel()
+	testutil.WaitFor(t, "queued cancellation", func() bool { return queued.State() == SessionFailed })
+	if !errors.Is(queued.Err(), context.Canceled) {
+		t.Fatalf("queued error = %v", queued.Err())
+	}
+	if s, err := mgr.Start(ctx, WorkerSpec{Session: "canceled"}); s != nil || !errors.Is(err, context.Canceled) {
+		t.Fatalf("canceled Start = (%v, %v)", s, err)
+	}
+	retry, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"})
+	if err != nil {
+		t.Fatalf("capacity not released after cancellation: %v", err)
+	}
+	release.Do(func() { close(block) })
+	testutil.WaitFor(t, "first completed", func() bool { return first.State() == SessionDone })
+	testutil.WaitFor(t, "retry failed", func() bool { return retry.State() == SessionFailed })
+	if !errors.Is(retry.Err(), ErrPermanent) {
+		t.Fatalf("retry error = %v", retry.Err())
+	}
+	last, err := mgr.Start(context.Background(), WorkerSpec{Session: "last"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	testutil.WaitFor(t, "last completed", func() bool { return last.State() == SessionDone })
+	mgr.mu.Lock()
+	defer mgr.mu.Unlock()
+	if mgr.pending != 0 || len(mgr.sessions) != 4 {
+		t.Fatalf("pending=%d retained=%d, want 0 and 4", mgr.pending, len(mgr.sessions))
+	}
+}
+
+func TestSessionManagerPrunesOnCompletion(t *testing.T) {
+	block := make(chan struct{})
+	var release sync.Once
+	defer release.Do(func() { close(block) })
+	launcher, _ := seqLauncher(
+		&fakeWorker{pid: 1, waitCh: block},
+		&fakeWorker{pid: 2},
+		&fakeWorker{pid: 3, exitCode: ExitPermanent},
+	)
+	pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
+	mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2})
+	sessions := make([]*Session, 3)
+	for i := range sessions {
+		var err error
+		sessions[i], err = mgr.Start(context.Background(), WorkerSpec{Session: fmt.Sprintf("s%d", i)})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	release.Do(func() { close(block) })
+	for _, s := range sessions {
+		select {
+		case <-s.Done():
+		case <-time.After(2 * time.Second):
+			t.Fatal("session did not finish")
+		}
+	}
+	if got := len(mgr.Statuses()); got != 2 {
+		t.Fatalf("retained %d completed sessions without another Start, want 2", got)
+	}
+}
+
 func TestSessionStatuses(t *testing.T) {
 	launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"l"}, exitCode: 0})
 	pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})