feat: manage task process lifetimes and preserve turn history

internal/runlease/scope.go

@@ -0,0 +1,126 @@
+// Package runlease binds asynchronous tool workers to a task run even when
+// their contexts detach from a per-call timeout or an SSE connection.
+package runlease
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+	"sync"
+)
+
+var ErrClosed = errors.New("task is ending; new tool executions are not allowed")
+var ErrUnconfirmed = errors.New("remote cancellation is unconfirmed")
+
+// Bound detached worker fan-out independently of OS process limits.
+const MaxTaskWorkers = 256
+
+type contextKey struct{}
+type Scope struct {
+	mu          sync.Mutex
+	sealed      bool
+	workers     map[string]context.CancelFunc
+	unconfirmed map[string]string
+	changed     chan struct{}
+}
+
+func New() *Scope {
+	return &Scope{workers: make(map[string]context.CancelFunc), unconfirmed: make(map[string]string), changed: make(chan struct{})}
+}
+func WithScope(ctx context.Context, s *Scope) context.Context {
+	return context.WithValue(ctx, contextKey{}, s)
+}
+func FromContext(ctx context.Context) *Scope {
+	if ctx == nil {
+		return nil
+	}
+	s, _ := ctx.Value(contextKey{}).(*Scope)
+	return s
+}
+func (s *Scope) notify() { close(s.changed); s.changed = make(chan struct{}) }
+func (s *Scope) Register(id string, cancel context.CancelFunc) (func(), error) {
+	if s == nil {
+		return func() {}, nil
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.sealed {
+		return nil, ErrClosed
+	}
+	if len(s.workers) >= MaxTaskWorkers {
+		return nil, fmt.Errorf("task worker limit reached (%d)", MaxTaskWorkers)
+	}
+	if _, ok := s.workers[id]; ok {
+		return nil, fmt.Errorf("duplicate worker %s", id)
+	}
+	s.workers[id] = cancel
+	var once sync.Once
+	return func() { once.Do(func() { s.mu.Lock(); delete(s.workers, id); s.notify(); s.mu.Unlock() }) }, nil
+}
+func (s *Scope) Seal() {
+	if s == nil {
+		return
+	}
+	s.mu.Lock()
+	s.sealed = true
+	s.mu.Unlock()
+}
+func (s *Scope) Cancel() {
+	if s == nil {
+		return
+	}
+	s.mu.Lock()
+	s.sealed = true
+	cs := make([]context.CancelFunc, 0, len(s.workers))
+	for _, c := range s.workers {
+		cs = append(cs, c)
+	}
+	s.mu.Unlock()
+	for _, c := range cs {
+		if c != nil {
+			c()
+		}
+	}
+}
+func (s *Scope) MarkUnconfirmed(id, message string) {
+	if s == nil {
+		return
+	}
+	s.mu.Lock()
+	s.unconfirmed[id] = message
+	s.notify()
+	s.mu.Unlock()
+}
+func (s *Scope) Wait(ctx context.Context) error {
+	if s == nil {
+		return nil
+	}
+	for {
+		s.mu.Lock()
+		pending := make([]string, 0, len(s.workers))
+		for id := range s.workers {
+			pending = append(pending, id)
+		}
+		uncertain := make([]string, 0, len(s.unconfirmed))
+		for id, msg := range s.unconfirmed {
+			uncertain = append(uncertain, id+": "+msg)
+		}
+		changed := s.changed
+		s.mu.Unlock()
+		if len(pending) == 0 {
+			if len(uncertain) > 0 {
+				sort.Strings(uncertain)
+				return fmt.Errorf("%w: %s", ErrUnconfirmed, strings.Join(uncertain, "; "))
+			}
+			return nil
+		}
+		select {
+		case <-changed:
+		case <-ctx.Done():
+			sort.Strings(pending)
+			return fmt.Errorf("tool workers still running %v: %w", pending, ctx.Err())
+		}
+	}
+}

internal/runlease/scope_test.go

@@ -0,0 +1,70 @@
+package runlease
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestConcurrentAdmissionAndCancellation(t *testing.T) {
+	scope := New()
+	ctx := WithScope(context.Background(), scope)
+	if FromContext(context.WithoutCancel(ctx)) != scope {
+		t.Fatal("detachment lost task ownership")
+	}
+	var wg sync.WaitGroup
+	for i := 0; i < 64; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			worker, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			release, err := scope.Register(string(rune('a'+i)), cancel)
+			if errors.Is(err, ErrClosed) {
+				return
+			}
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			<-worker.Done()
+			release()
+		}(i)
+	}
+	scope.Cancel()
+	wg.Wait()
+	wait, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := scope.Wait(wait); err != nil {
+		t.Fatal(err)
+	}
+}
+func TestDetachedWorkerCapacityReleasedOnCompletion(t *testing.T) {
+	scope := New()
+	releases := make([]func(), 0, MaxTaskWorkers)
+	for i := 0; i < MaxTaskWorkers; i++ {
+		release, err := scope.Register(string(rune(i)), func() {})
+		if err != nil {
+			t.Fatal(err)
+		}
+		releases = append(releases, release)
+	}
+	if _, err := scope.Register("overflow", func() {}); err == nil {
+		t.Fatal("unbounded worker admission")
+	}
+	releases[0]()
+	release, err := scope.Register("replacement", func() {})
+	if err != nil {
+		t.Fatal(err)
+	}
+	release()
+	for _, release := range releases {
+		release()
+	}
+	scope.Cancel()
+	if err = scope.Wait(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+}