feat(helix-org): session-layer transcript mirror + fire-and-forget worker sessions

api/pkg/org/application/lifecycle/lifecycle.go

@@ -51,6 +51,10 @@ type Service struct {
 	// collapses an ex-manager's team Stream when its last report just
 	// left. nil is a no-op (tests without topology wiring).
 	Topology *topology.Reconciler
+
+	// Mirror is the transcript mirror; Fire stops the fired Worker's
+	// subscription so it doesn't leak. nil is a no-op.
+	Mirror *helix.Mirror
 }
 
 // ErrOwnerProtected is returned by Fire when the caller targets the
@@ -116,6 +120,10 @@ func (s *Service) Fire(ctx context.Context, orgID string, id orgchart.WorkerID)
 
 	state, _ := helix.LoadState(ctx, s.Store, orgID, id)
 
+	if s.Mirror != nil {
+		s.Mirror.Stop(id)
+	}
+
 	if s.Helix != nil && state.ProjectID != "" {
 		if err := s.Helix.DeleteProject(ctx, state.ProjectID); err != nil && !errors.Is(err, helix.ErrProjectNotFound) {
 			s.logger().Warn("fire: delete helix project", "worker", id, "project", state.ProjectID, "err", err)

api/pkg/org/infrastructure/runtime/helix/mirror.go

@@ -0,0 +1,241 @@
+package helix
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/helixml/helix/api/pkg/org/application/agent"
+	"github.com/helixml/helix/api/pkg/org/application/streamhub"
+	"github.com/helixml/helix/api/pkg/org/domain/orgchart"
+	"github.com/helixml/helix/api/pkg/org/domain/store"
+	"github.com/helixml/helix/api/pkg/pubsub"
+	"github.com/helixml/helix/api/pkg/types"
+)
+
+const defaultMirrorPoll = 5 * time.Second
+
+// MirrorConfig wires the session-layer transcript Mirror.
+type MirrorConfig struct {
+	PubSub      pubsub.PubSub
+	Snapshotter SessionPreamble
+	// Client resolves the session owner so we subscribe to the right
+	// GetSessionQueue(owner, session) topic.
+	Client SpawnerClient
+	// ExploratorySession returns a project's current exploratory session
+	// (the one the inline chat / live UI follow). The mirror polls it to
+	// track the worker as its session churns. "" means no session yet.
+	ExploratorySession func(ctx context.Context, projectID string) (string, error)
+	Store              *store.Store
+	Hub                *streamhub.Hub
+	NewID              func() string
+	Now                func() time.Time
+	Logger             *slog.Logger
+	PollInterval       time.Duration // <=0 uses defaultMirrorPoll; seam for tests
+}
+
+// Mirror is the single writer of worker transcript segments: it keeps
+// one subscription per tracked worker pointed at that worker's current
+// session and republishes every settled entry (plus the user prompt)
+// onto s-activations-<worker>. Every turn — spawner activation, inline
+// chat, anything on /sessions/chat — flows through the session topic, so
+// one subscriber captures them all.
+//
+// A worker's session is not stable (stale resume opens a fresh one;
+// inline chat can land on a newer one than the spawner persisted), so we
+// track the worker — whose project is stable — and poll its current
+// exploratory session, re-pointing when it changes.
+type Mirror struct {
+	base context.Context
+	cfg  MirrorConfig
+
+	mu      sync.Mutex
+	tracked map[orgchart.WorkerID]context.CancelFunc
+}
+
+// NewMirror constructs a Mirror. base bounds every tracker (typically
+// the server lifetime). A nil PubSub disables the mirror.
+func NewMirror(base context.Context, cfg MirrorConfig) *Mirror {
+	if base == nil {
+		base = context.Background()
+	}
+	return &Mirror{
+		base:    base,
+		cfg:     cfg,
+		tracked: map[orgchart.WorkerID]context.CancelFunc{},
+	}
+}
+
+func (m *Mirror) pollInterval() time.Duration {
+	if m.cfg.PollInterval > 0 {
+		return m.cfg.PollInterval
+	}
+	return defaultMirrorPoll
+}
+
+// Ensure starts tracking a worker (idempotent). The session is resolved
+// by the tracker, not passed in — the caller's notion of "the session"
+// is exactly what goes stale.
+func (m *Mirror) Ensure(orgID string, workerID orgchart.WorkerID) {
+	if m == nil || m.cfg.PubSub == nil {
+		return
+	}
+	m.mu.Lock()
+	if _, ok := m.tracked[workerID]; ok {
+		m.mu.Unlock()
+		return
+	}
+	ctx, cancel := context.WithCancel(m.base)
+	m.tracked[workerID] = cancel
+	m.mu.Unlock()
+	go m.track(ctx, orgID, workerID)
+}
+
+// EnsureAll tracks every worker in an org — called from bootstrap so
+// pre-existing / inline-chat-only workers are mirrored without an
+// activation first.
+func (m *Mirror) EnsureAll(ctx context.Context, orgID string) {
+	if m == nil || m.cfg.PubSub == nil || m.cfg.Store == nil {
+		return
+	}
+	workers, err := m.cfg.Store.Workers.List(ctx, orgID)
+	if err != nil {
+		if m.cfg.Logger != nil {
+			m.cfg.Logger.Warn("helix mirror: list workers for sweep", "org", orgID, "err", err)
+		}
+		return
+	}
+	for _, w := range workers {
+		m.Ensure(orgID, w.ID())
+	}
+}
+
+// Stop stops tracking a worker (on fire).
+func (m *Mirror) Stop(workerID orgchart.WorkerID) {
+	if m == nil {
+		return
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if cancel, ok := m.tracked[workerID]; ok {
+		cancel()
+		delete(m.tracked, workerID)
+	}
+}
+
+// track resolves the worker's current session and (re)subscribes
+// whenever it changes, polling on an interval, until ctx fires. Each
+// session gets a fresh bridge.
+func (m *Mirror) track(ctx context.Context, orgID string, workerID orgchart.WorkerID) {
+	var (
+		curSession string
+		curCancel  = func() {}
+	)
+	defer func() { curCancel() }()
+
+	repoint := func() {
+		desired := m.resolveSession(ctx, orgID, workerID)
+		if desired == "" || desired == curSession {
+			return
+		}
+		curCancel() // old pump flushes pending entries on ctx.Done
+		subCtx, cancel := context.WithCancel(ctx)
+		curCancel = cancel
+		curSession = desired
+		m.subscribe(subCtx, orgID, workerID, desired)
+		if m.cfg.Logger != nil {
+			m.cfg.Logger.Info("helix mirror: pointed at session", "worker", workerID, "session", desired)
+		}
+	}
+
+	repoint() // immediately, don't wait for the first tick
+	ticker := time.NewTicker(m.pollInterval())
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			repoint()
+		}
+	}
+}
+
+// resolveSession returns the worker's current exploratory session,
+// falling back to the persisted pointer when that lookup is unavailable.
+func (m *Mirror) resolveSession(ctx context.Context, orgID string, workerID orgchart.WorkerID) string {
+	state, err := LoadState(ctx, m.cfg.Store, orgID, workerID)
+	if err != nil {
+		return ""
+	}
+	if state.ProjectID != "" && m.cfg.ExploratorySession != nil {
+		if sid, err := m.cfg.ExploratorySession(ctx, state.ProjectID); err == nil && sid != "" {
+			return sid
+		} else if err != nil && m.cfg.Logger != nil {
+			m.cfg.Logger.Warn("helix mirror: resolve exploratory session", "worker", workerID, "project", state.ProjectID, "err", err)
+		}
+	}
+	return state.SessionID
+}
+
+// subscribe attaches a bridge to one session's topic and pumps its
+// frames onto the activation stream until ctx fires. The first subscribe
+// is synchronous so a frame published right after isn't raced.
+func (m *Mirror) subscribe(ctx context.Context, orgID string, workerID orgchart.WorkerID, sessionID string) {
+	ownerID := ""
+	if m.cfg.Client != nil {
+		if owner, err := m.cfg.Client.SessionOwner(ctx, sessionID); err != nil {
+			if m.cfg.Logger != nil {
+				m.cfg.Logger.Warn("helix mirror: resolve session owner", "worker", workerID, "session", sessionID, "err", err)
+			}
+		} else {
+			ownerID = owner
+		}
+	}
+	publish := func(body string) {
+		if body == "" {
+			return
+		}
+		_, _ = agent.PublishActivationEvent(ctx, m.cfg.Store, m.cfg.Hub, m.cfg.NewID, m.cfg.Now, m.cfg.Logger, orgID, workerID, body)
+	}
+	b := newBridge(publish)
+	ch, err := SubscribeSessionUpdates(ctx, m.cfg.PubSub, m.cfg.Snapshotter, ownerID, sessionID)
+	if err != nil {
+		if m.cfg.Logger != nil {
+			m.cfg.Logger.Warn("helix mirror: subscribe", "worker", workerID, "session", sessionID, "err", err)
+		}
+		ch = nil
+	}
+	go m.pump(ctx, b, ownerID, sessionID, ch)
+}
+
+// pump drains the subscription channel into the bridge and reconnects
+// with capped backoff until ctx fires.
+func (m *Mirror) pump(ctx context.Context, b *bridge, ownerID, sessionID string, ch <-chan types.WebsocketEvent) {
+	delay := time.Second
+	for {
+		if ch != nil {
+			for u := range ch {
+				b.apply(u)
+			}
+		}
+		select {
+		case <-ctx.Done():
+			b.stream.Flush()
+			return
+		case <-time.After(delay):
+		}
+		if delay < 30*time.Second {
+			delay *= 2
+		}
+		var err error
+		ch, err = SubscribeSessionUpdates(ctx, m.cfg.PubSub, m.cfg.Snapshotter, ownerID, sessionID)
+		if err != nil {
+			if m.cfg.Logger != nil {
+				m.cfg.Logger.Warn("helix mirror: re-subscribe", "session", sessionID, "err", err)
+			}
+			ch = nil
+		}
+	}
+}