kandev/apps/backend/internal/agent/runtime/lifecycle/stream_coalescer.go

package lifecycle

import (
	"sync"
	"time"
)

const defaultStreamCoalesceWindow = 100 * time.Millisecond

type coalescedStreamChunk struct {
	eventType        string
	messageID        string
	attemptID        string
	content          string
	isAppend         bool
	diagnostic       bool
	promptGeneration uint64
}

// streamCoalescer combines adjacent append chunks for one execution. The
// first non-empty chunk of a record is emitted immediately; later chunks wait for the
// bounded window or an explicit lifecycle boundary. A single pending segment
// is intentional: combining across another message ID would change wire
// ordering.
type streamCoalescer struct {
	emitMu               sync.Mutex
	mu                   sync.Mutex
	window               time.Duration
	pending              *coalescedStreamChunk
	timer                *time.Timer
	closed               bool
	publish              func(coalescedStreamChunk)
	lastEventType        string
	lastMessageID        string
	lastAttemptID        string
	lastDiagnostic       bool
	lastPromptGeneration uint64
	forceImmediate       bool
	received             int
	coalesced            int
	flushed              int
}

type streamCoalescerStats struct {
	received  int
	coalesced int
	flushed   int
}

func newStreamCoalescer(window time.Duration, publish func(coalescedStreamChunk)) *streamCoalescer {
	if window <= 0 {
		window = defaultStreamCoalesceWindow
	}
	return &streamCoalescer{window: window, publish: publish}
}

func (c *streamCoalescer) add(chunk coalescedStreamChunk) {
	if chunk.content == "" {
		return
	}

	c.emitMu.Lock()
	defer c.emitMu.Unlock()

	var ready []coalescedStreamChunk
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return
	}
	c.received++

	// A diagnostic-marker change is treated as a correlation-key change, just
	// like a messageID or eventType change: merging a marked chunk's content
	// into an unmarked pending segment (or vice versa) would silently erase
	// the marker for the merged text. A promptGeneration change gets the same
	// treatment: merging chunks from two different prompt attempts would stamp
	// the merged text with only one attempt's generation, breaking downstream
	// recovery-evidence correlation for the other attempt's content. An
	// attemptID change is likewise a correlation-key change: it identifies the
	// immutable recovery attempt that owns the callback, so merging across an
	// attemptID boundary would relabel one attempt's content with another's.
	sameAsLast := c.lastEventType == chunk.eventType && c.lastMessageID == chunk.messageID &&
		c.lastAttemptID == chunk.attemptID &&
		c.lastDiagnostic == chunk.diagnostic && c.lastPromptGeneration == chunk.promptGeneration
	immediate := !chunk.isAppend || c.forceImmediate || !sameAsLast
	c.forceImmediate = false
	switch {
	case immediate:
		ready = c.detachLocked(ready)
		ready = append(ready, chunk)
	case c.pending != nil && c.pending.eventType == chunk.eventType && c.pending.messageID == chunk.messageID &&
		c.pending.attemptID == chunk.attemptID &&
		c.pending.diagnostic == chunk.diagnostic && c.pending.promptGeneration == chunk.promptGeneration:
		c.pending.content += chunk.content
		c.coalesced++
	default:
		ready = c.detachLocked(ready)
		pending := chunk
		c.pending = &pending
		c.coalesced++
		c.timer = time.AfterFunc(c.window, c.flush)
	}
	c.lastEventType = chunk.eventType
	c.lastMessageID = chunk.messageID
	c.lastAttemptID = chunk.attemptID
	c.lastDiagnostic = chunk.diagnostic
	c.lastPromptGeneration = chunk.promptGeneration
	c.mu.Unlock()

	c.publishReady(ready)
}

func (c *streamCoalescer) flush() {
	c.emitMu.Lock()
	defer c.emitMu.Unlock()
	c.mu.Lock()
	ready := c.detachLocked(nil)
	c.mu.Unlock()
	c.publishReady(ready)
}

func (c *streamCoalescer) flushBoundary() {
	c.emitMu.Lock()
	defer c.emitMu.Unlock()
	c.mu.Lock()
	c.forceImmediate = true
	ready := c.detachLocked(nil)
	c.mu.Unlock()
	c.publishReady(ready)
}

func (c *streamCoalescer) close() {
	c.emitMu.Lock()
	defer c.emitMu.Unlock()
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return
	}
	c.closed = true
	ready := c.detachLocked(nil)
	c.mu.Unlock()
	c.publishReady(ready)
}

func (c *streamCoalescer) detachLocked(ready []coalescedStreamChunk) []coalescedStreamChunk {
	// A timer callback that already fired can be waiting on emitMu while add
	// installs a later pending segment and timer. When that callback acquires
	// emitMu it may detach the later segment immediately, collapsing its window;
	// content and ordering remain correct because all detaches are serialized.
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	if c.pending != nil {
		ready = append(ready, *c.pending)
		c.pending = nil
		c.flushed++
	}
	return ready
}

func (c *streamCoalescer) stats() streamCoalescerStats {
	c.mu.Lock()
	defer c.mu.Unlock()
	return streamCoalescerStats{
		received:  c.received,
		coalesced: c.coalesced,
		flushed:   c.flushed,
	}
}

func (c *streamCoalescer) publishReady(chunks []coalescedStreamChunk) {
	if c.publish == nil {
		return
	}
	for _, chunk := range chunks {
		c.publish(chunk)
	}
}