fix: chatd refactor (#26270)

Implements the chatd stabilization RFC. Combines: - https://github.com/coder/coder/pull/25908 - https://github.com/coder/coder/pull/25923 - https://github.com/coder/coder/pull/26109 - https://github.com/coder/coder/pull/26110 - https://github.com/coder/coder/pull/26111 - https://github.com/coder/coder/pull/26112
2026-09-24 15:04:27 +08:00 · 2026-06-12 13:33:12 +02:00
parent 4a07f61c50
commit 4debd23cbb
155 changed files with 37612 additions and 24513 deletions
@@ -0,0 +1,9 @@
+package messagepartbuffer
+
+// EpisodeCount returns the number of tracked episodes so tests can assert
+// that episode state is reclaimed and does not leak.
+func (b *Buffer) EpisodeCount() int {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return len(b.episodes)
+}
@@ -0,0 +1,494 @@
+package messagepartbuffer
+
+import (
+	"container/heap"
+	"context"
+	"encoding/json"
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/quartz"
+)
+
+const (
+	defaultMaxEpisodeBytes        = int64(1024 * 1024)
+	defaultClosedEpisodeRetention = 15 * time.Second
+	defaultSubscriberSendTimeout  = 10 * time.Second
+	defaultSubscriberChannelSize  = 16
+)
+
+var (
+	// ErrEpisodeExists means the episode already exists.
+	ErrEpisodeExists = xerrors.New("message part episode already exists")
+	// ErrEpisodeNotFound means the episode has not been created.
+	ErrEpisodeNotFound = xerrors.New("message part episode not found")
+	// ErrEpisodeClosed means the episode no longer accepts parts.
+	ErrEpisodeClosed = xerrors.New("message part episode closed")
+	// ErrEpisodeFull means the episode byte limit would be exceeded.
+	ErrEpisodeFull = xerrors.New("message part episode full")
+	// ErrMessagePartBufferClosed means the whole buffer is closed.
+	ErrMessagePartBufferClosed = xerrors.New("message part buffer closed")
+)
+
+// Key identifies a buffered message part episode.
+type Key struct {
+	ChatID            uuid.UUID
+	HistoryVersion    int64
+	GenerationAttempt int64
+}
+
+// Part is a buffered chat message part with its sequence number.
+type Part struct {
+	Seq         int64
+	Role        codersdk.ChatMessageRole
+	MessagePart codersdk.ChatMessagePart
+}
+
+type partJSON struct {
+	Seq  int64                    `json:"seq"`
+	Role codersdk.ChatMessageRole `json:"role"`
+	Part codersdk.ChatMessagePart `json:"part"`
+}
+
+func (p Part) jsonValue() partJSON {
+	return partJSON{
+		Seq:  p.Seq,
+		Role: p.Role,
+		Part: p.MessagePart,
+	}
+}
+
+// Options configures a Buffer.
+type Options struct {
+	MaxEpisodeBytes        int64
+	ClosedEpisodeRetention time.Duration
+	SubscriberSendTimeout  time.Duration
+	SubscriberChannelSize  int
+	Clock                  quartz.Clock
+}
+
+// Buffer stores streamed message parts by episode.
+type Buffer struct {
+	mu             sync.Mutex
+	opts           Options
+	episodes       map[Key]*episodeState
+	closedEpisodes closedEpisodeHeap
+	closed         bool
+	done           chan struct{}
+}
+
+type episodeState struct {
+	created        bool
+	closed         bool
+	closedAt       time.Time
+	closedHeapItem *closedEpisodeItem
+	parts          []Part
+	bytes          int64
+	subscribers    map[*episodeSubscriber]struct{}
+}
+
+type closedEpisodeItem struct {
+	key      Key
+	closedAt time.Time
+}
+
+type closedEpisodeHeap []*closedEpisodeItem
+
+func (h closedEpisodeHeap) Len() int {
+	return len(h)
+}
+
+func (h closedEpisodeHeap) Less(i, j int) bool {
+	return h[i].closedAt.Before(h[j].closedAt)
+}
+
+func (h closedEpisodeHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+}
+
+func (h *closedEpisodeHeap) Push(value any) {
+	item, ok := value.(*closedEpisodeItem)
+	if !ok {
+		// The reason we panic here instead of returning an error is that
+		// closedEpisodeHeap implements the https://pkg.go.dev/container/heap interface.
+		// We must accept an any type and we must not return an error.
+		panic("closed episode heap received invalid item")
+	}
+	*h = append(*h, item)
+}
+
+func (h *closedEpisodeHeap) Pop() any {
+	old := *h
+	last := old[len(old)-1]
+	old[len(old)-1] = nil
+	*h = old[:len(old)-1]
+	return last
+}
+
+type episodeSubscriber struct {
+	out      chan Part
+	notifyCh chan struct{}
+	stopCh   chan struct{}
+	next     int
+	stopOnce sync.Once
+}
+
+// New returns a message part buffer.
+func New(options Options) *Buffer {
+	if options.MaxEpisodeBytes <= 0 {
+		options.MaxEpisodeBytes = defaultMaxEpisodeBytes
+	}
+	if options.ClosedEpisodeRetention <= 0 {
+		options.ClosedEpisodeRetention = defaultClosedEpisodeRetention
+	}
+	if options.SubscriberSendTimeout <= 0 {
+		options.SubscriberSendTimeout = defaultSubscriberSendTimeout
+	}
+	if options.SubscriberChannelSize <= 0 {
+		options.SubscriberChannelSize = defaultSubscriberChannelSize
+	}
+	if options.Clock == nil {
+		options.Clock = quartz.NewReal()
+	}
+	buffer := &Buffer{
+		opts:     options,
+		episodes: make(map[Key]*episodeState),
+		done:     make(chan struct{}),
+	}
+	buffer.startCleanupLoop()
+	return buffer
+}
+
+// CreateEpisode creates a new episode.
+func (b *Buffer) CreateEpisode(key Key) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closed {
+		return ErrMessagePartBufferClosed
+	}
+	b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "create"))
+	episode := b.episodeLocked(key)
+	if episode.created {
+		return ErrEpisodeExists
+	}
+	episode.created = true
+	return nil
+}
+
+// AddPart appends a part to an existing episode.
+func (b *Buffer) AddPart(key Key, role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closed {
+		return ErrMessagePartBufferClosed
+	}
+	episode := b.episodes[key]
+	if episode == nil || !episode.created {
+		return ErrEpisodeNotFound
+	}
+	if episode.closed {
+		return ErrEpisodeClosed
+	}
+	buffered := Part{
+		Seq:         int64(len(episode.parts) + 1),
+		Role:        role,
+		MessagePart: part,
+	}
+	sizeBytes, err := serializedPartBytes(buffered)
+	if err != nil {
+		return err
+	}
+	if episode.bytes+sizeBytes > b.opts.MaxEpisodeBytes {
+		return ErrEpisodeFull
+	}
+	episode.parts = append(episode.parts, buffered)
+	episode.bytes += sizeBytes
+	for subscriber := range episode.subscribers {
+		notifySubscriber(subscriber)
+	}
+	return nil
+}
+
+// CloseEpisode marks an episode closed and closes its subscribers.
+func (b *Buffer) CloseEpisode(key Key) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closed {
+		return ErrMessagePartBufferClosed
+	}
+	episode := b.episodeLocked(key)
+	episode.created = true
+	if episode.closed {
+		return nil
+	}
+	episode.closed = true
+	episode.closedAt = b.opts.Clock.Now("message-part-buffer", "close")
+	b.queueClosedEpisodeLocked(key, episode)
+	for subscriber := range episode.subscribers {
+		notifySubscriber(subscriber)
+	}
+	return nil
+}
+
+// GetParts returns a snapshot of buffered parts for an episode.
+func (b *Buffer) GetParts(key Key) ([]Part, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closed {
+		return nil, ErrMessagePartBufferClosed
+	}
+	b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "get"))
+	episode := b.episodes[key]
+	if episode == nil || !episode.created {
+		return nil, ErrEpisodeNotFound
+	}
+	return append([]Part(nil), episode.parts...), nil
+}
+
+// SubscribeToEpisode replays existing parts and streams new parts.
+func (b *Buffer) SubscribeToEpisode(ctx context.Context, key Key) (<-chan Part, func(), error) {
+	b.mu.Lock()
+	if b.closed {
+		b.mu.Unlock()
+		return nil, nil, ErrMessagePartBufferClosed
+	}
+	episode := b.episodeLocked(key)
+	subscriber := &episodeSubscriber{
+		out:      make(chan Part),
+		notifyCh: make(chan struct{}, 1),
+		stopCh:   make(chan struct{}),
+	}
+	if episode.subscribers == nil {
+		episode.subscribers = make(map[*episodeSubscriber]struct{})
+	}
+	episode.subscribers[subscriber] = struct{}{}
+	notifySubscriber(subscriber)
+	b.mu.Unlock()
+
+	go b.deliverSubscriber(ctx, key, subscriber)
+	cancel := func() {
+		b.cancelSubscriber(key, subscriber)
+	}
+	return subscriber.out, cancel, nil
+}
+
+// Close closes the buffer and all pending subscriptions.
+func (b *Buffer) Close() {
+	b.mu.Lock()
+	if b.closed {
+		b.mu.Unlock()
+		return
+	}
+	b.closed = true
+	close(b.done)
+	for _, episode := range b.episodes {
+		for subscriber := range episode.subscribers {
+			b.stopSubscriberLocked(episode, subscriber)
+		}
+	}
+	b.mu.Unlock()
+}
+
+func (b *Buffer) startCleanupLoop() {
+	ticker := b.opts.Clock.NewTicker(b.opts.ClosedEpisodeRetention, "message-part-buffer", "cleanup")
+	go func() {
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				b.mu.Lock()
+				if b.closed {
+					b.mu.Unlock()
+					return
+				}
+				b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "cleanup"))
+				b.mu.Unlock()
+			case <-b.done:
+				return
+			}
+		}
+	}()
+}
+
+func (b *Buffer) gcClosedEpisodesLocked(now time.Time) {
+	cutoff := now.Add(-b.opts.ClosedEpisodeRetention)
+	type retainedEpisode struct {
+		key     Key
+		episode *episodeState
+	}
+	retained := make([]retainedEpisode, 0)
+	for b.closedEpisodes.Len() > 0 {
+		item := b.closedEpisodes[0]
+		if item.closedAt.After(cutoff) {
+			break
+		}
+		popped, ok := heap.Pop(&b.closedEpisodes).(*closedEpisodeItem)
+		if !ok || popped != item {
+			continue
+		}
+		episode := b.episodes[item.key]
+		if episode == nil || episode.closedHeapItem != item || !episode.closed {
+			continue
+		}
+		episode.closedHeapItem = nil
+		if len(episode.subscribers) > 0 {
+			retained = append(retained, retainedEpisode{key: item.key, episode: episode})
+			continue
+		}
+		delete(b.episodes, item.key)
+	}
+	for _, item := range retained {
+		if b.episodes[item.key] != item.episode || !item.episode.closed || item.episode.closedHeapItem != nil {
+			continue
+		}
+		b.queueClosedEpisodeLocked(item.key, item.episode)
+	}
+}
+
+func (b *Buffer) queueClosedEpisodeLocked(key Key, episode *episodeState) {
+	if episode.closedHeapItem != nil {
+		return
+	}
+	item := &closedEpisodeItem{key: key, closedAt: episode.closedAt}
+	episode.closedHeapItem = item
+	heap.Push(&b.closedEpisodes, item)
+}
+
+func (b *Buffer) episodeLocked(key Key) *episodeState {
+	episode := b.episodes[key]
+	if episode != nil {
+		return episode
+	}
+	episode = &episodeState{}
+	b.episodes[key] = episode
+	return episode
+}
+
+func (b *Buffer) subscriberParts(key Key, subscriber *episodeSubscriber) (parts []Part, closed bool, ok bool) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closed {
+		return nil, false, false
+	}
+	episode := b.episodes[key]
+	if episode == nil {
+		return nil, false, false
+	}
+	if !episode.created {
+		return nil, false, true
+	}
+	if subscriber.next > len(episode.parts) {
+		return nil, false, false
+	}
+	parts = append([]Part(nil), episode.parts[subscriber.next:]...)
+	subscriber.next = len(episode.parts)
+	return parts, episode.closed && subscriber.next == len(episode.parts), true
+}
+
+func (b *Buffer) deliverSubscriber(ctx context.Context, key Key, subscriber *episodeSubscriber) {
+	defer close(subscriber.out)
+	defer b.removeSubscriber(key, subscriber)
+	for {
+		parts, closed, ok := b.subscriberParts(key, subscriber)
+		if !ok {
+			return
+		}
+		for _, part := range parts {
+			if !b.sendSubscriberPart(ctx, subscriber, part) {
+				return
+			}
+		}
+		if closed {
+			return
+		}
+		select {
+		case <-subscriber.notifyCh:
+		case <-subscriber.stopCh:
+			return
+		case <-ctx.Done():
+			return
+		case <-b.done:
+			return
+		}
+	}
+}
+
+func (b *Buffer) sendSubscriberPart(ctx context.Context, subscriber *episodeSubscriber, part Part) bool {
+	timer := b.opts.Clock.NewTimer(b.opts.SubscriberSendTimeout, "message-part-buffer", "subscriber-send")
+	defer timer.Stop()
+	select {
+	case subscriber.out <- part:
+		return true
+	case <-timer.C:
+		return false
+	case <-subscriber.stopCh:
+		return false
+	case <-ctx.Done():
+		return false
+	case <-b.done:
+		return false
+	}
+}
+
+func (b *Buffer) cancelSubscriber(key Key, subscriber *episodeSubscriber) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	episode := b.episodes[key]
+	if episode != nil {
+		b.stopSubscriberLocked(episode, subscriber)
+		return
+	}
+	subscriber.stop()
+}
+
+func (b *Buffer) removeSubscriber(key Key, subscriber *episodeSubscriber) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	episode := b.episodes[key]
+	if episode == nil {
+		return
+	}
+	delete(episode.subscribers, subscriber)
+	if len(episode.subscribers) != 0 {
+		return
+	}
+	switch {
+	case episode.closed:
+		b.queueClosedEpisodeLocked(key, episode)
+	case !episode.created:
+		// SubscribeToEpisode inserts a placeholder state for unknown keys so
+		// that CreateEpisode can adopt subscribers that arrive early. Once the
+		// last subscriber leaves a still-uncreated episode, no CreateEpisode or
+		// CloseEpisode call will ever reclaim it, so delete it here to avoid
+		// leaking the map entry for the lifetime of the buffer.
+		delete(b.episodes, key)
+	}
+}
+
+func (*Buffer) stopSubscriberLocked(episode *episodeState, subscriber *episodeSubscriber) {
+	delete(episode.subscribers, subscriber)
+	subscriber.stop()
+}
+
+func notifySubscriber(subscriber *episodeSubscriber) {
+	select {
+	case subscriber.notifyCh <- struct{}{}:
+	default:
+	}
+}
+
+func (s *episodeSubscriber) stop() {
+	s.stopOnce.Do(func() { close(s.stopCh) })
+}
+
+func serializedPartBytes(part Part) (int64, error) {
+	data, err := json.Marshal(part.jsonValue())
+	if err != nil {
+		return 0, err
+	}
+	return int64(len(data)), nil
+}
@@ -0,0 +1,395 @@
+package messagepartbuffer_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
+)
+
+func TestBuffer_CreateEpisodeRejectsDuplicate(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.ErrorIs(t, buffer.CreateEpisode(key), messagepartbuffer.ErrEpisodeExists)
+}
+
+func TestBuffer_AddPartAndGetParts(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello")))
+
+	parts, err := buffer.GetParts(key)
+	require.NoError(t, err)
+	require.Len(t, parts, 1)
+	require.Equal(t, int64(1), parts[0].Seq)
+	require.Equal(t, codersdk.ChatMessageRoleAssistant, parts[0].Role)
+	require.Equal(t, codersdk.ChatMessageText("hello"), parts[0].MessagePart)
+}
+
+func TestBuffer_AddPartMissingEpisodeReturnsNotFound(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	err := buffer.AddPart(testEpisodeKey(), codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello"))
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+}
+
+func TestBuffer_GetPartsMissingEpisodeReturnsNotFound(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	_, err := buffer.GetParts(testEpisodeKey())
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+}
+
+func TestBuffer_AddPartFullEpisodeReturnsFull(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{MaxEpisodeBytes: 1})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	err := buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello"))
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeFull)
+	parts, getErr := buffer.GetParts(key)
+	require.NoError(t, getErr)
+	require.Empty(t, parts)
+}
+
+func TestBuffer_CloseEpisodeMissingCreatesClosedEpisode(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CloseEpisode(key))
+	parts, err := buffer.GetParts(key)
+	require.NoError(t, err)
+	require.Empty(t, parts)
+	err = buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("tail"))
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeClosed)
+}
+
+func TestBuffer_CloseEpisodeIdempotent(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.CloseEpisode(key))
+	require.NoError(t, buffer.CloseEpisode(key))
+}
+
+func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("before")))
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+	require.Equal(t, "before", receivePart(t, ch).MessagePart.Text)
+
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("after")))
+	require.Equal(t, "after", receivePart(t, ch).MessagePart.Text)
+}
+
+func TestBuffer_SubscribeClosedEpisodeReplaysThenCloses(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("before")))
+	require.NoError(t, buffer.CloseEpisode(key))
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+	require.Equal(t, "before", receivePart(t, ch).MessagePart.Text)
+	assertChannelClosed(t, ch)
+}
+
+func TestBuffer_SubscribeBeforeCreateReturnsAndWaitsWithoutNotFound(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+
+	select {
+	case part := <-ch:
+		t.Fatalf("received part before episode create: %+v", part)
+	default:
+	}
+
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("live")))
+	require.Equal(t, "live", receivePart(t, ch).MessagePart.Text)
+}
+
+func TestBuffer_AddPartAssignsContiguousSeq(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	for i := range 3 {
+		require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText(string(rune('a'+i)))))
+	}
+	parts, err := buffer.GetParts(key)
+	require.NoError(t, err)
+	require.Equal(t, []int64{1, 2, 3}, []int64{parts[0].Seq, parts[1].Seq, parts[2].Seq})
+}
+
+func TestBuffer_EpisodeByteLimitUsesJSONAccounting(t *testing.T) {
+	t.Parallel()
+
+	part := codersdk.ChatMessageText("hello")
+	limit := serializedPartBytes(t, messagepartbuffer.Part{Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: part})
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{MaxEpisodeBytes: limit})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, part))
+	err := buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("too much"))
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeFull)
+}
+
+func TestBuffer_GCClosedEpisodeAfterGraceAndNoSubscribers(t *testing.T) {
+	t.Parallel()
+
+	clock := quartz.NewMock(t)
+	trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send")
+	defer trap.Close()
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{
+		Clock:                  clock,
+		ClosedEpisodeRetention: time.Minute,
+		SubscriberSendTimeout:  10 * time.Minute,
+	})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("held")))
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	require.NoError(t, buffer.CloseEpisode(key))
+	call := trap.MustWait(ctx)
+	call.MustRelease(ctx)
+	clock.Advance(time.Minute).MustWait(ctx)
+	clock.Advance(time.Second).MustWait(ctx)
+	_, err = buffer.GetParts(key)
+	require.NoError(t, err)
+
+	cancel()
+	drainUntilClosed(t, ch)
+	_, err = buffer.GetParts(key)
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+}
+
+func TestBuffer_GCRetainedSubscribedEpisodeDoesNotBlockOtherExpiredEpisodes(t *testing.T) {
+	t.Parallel()
+
+	clock := quartz.NewMock(t)
+	trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send")
+	defer trap.Close()
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{
+		Clock:                  clock,
+		ClosedEpisodeRetention: time.Minute,
+		SubscriberSendTimeout:  10 * time.Minute,
+	})
+	retainedKey := testEpisodeKey()
+	collectedKey := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(retainedKey))
+	require.NoError(t, buffer.AddPart(retainedKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("held")))
+	require.NoError(t, buffer.CreateEpisode(collectedKey))
+	require.NoError(t, buffer.AddPart(collectedKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("collect me")))
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, retainedKey)
+	require.NoError(t, err)
+	defer cancel()
+	require.NoError(t, buffer.CloseEpisode(retainedKey))
+	require.NoError(t, buffer.CloseEpisode(collectedKey))
+	call := trap.MustWait(ctx)
+	call.MustRelease(ctx)
+	clock.Advance(time.Minute).MustWait(ctx)
+	clock.Advance(time.Second).MustWait(ctx)
+
+	_, err = buffer.GetParts(retainedKey)
+	require.NoError(t, err)
+	_, err = buffer.GetParts(collectedKey)
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+
+	cancel()
+	drainUntilClosed(t, ch)
+	_, err = buffer.GetParts(retainedKey)
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+}
+
+func TestBuffer_SlowSubscriberClosed(t *testing.T) {
+	t.Parallel()
+
+	clock := quartz.NewMock(t)
+	trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send")
+	defer trap.Close()
+	stopTrap := clock.Trap().TimerStop()
+	defer stopTrap.Close()
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{
+		Clock:                 clock,
+		SubscriberSendTimeout: time.Second,
+	})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+
+	require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("blocked")))
+	call := trap.MustWait(ctx)
+	call.MustRelease(ctx)
+	clock.Advance(time.Second).MustWait(ctx)
+	stopCall := stopTrap.MustWait(ctx)
+	stopCall.MustRelease(ctx)
+	assertChannelClosed(t, ch)
+}
+
+func TestBuffer_BurstyOutputDoesNotCloseSubscriberBeforeSendTimeout(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{SubscriberChannelSize: 1})
+	key := testEpisodeKey()
+	require.NoError(t, buffer.CreateEpisode(key))
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+
+	for i := range 8 {
+		require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText(string(rune('a'+i)))))
+	}
+	for i := range 8 {
+		part := receivePart(t, ch)
+		require.Equal(t, string(rune('a'+i)), part.MessagePart.Text)
+	}
+}
+
+func TestBuffer_SubscribeCanceledBeforeCreateCanCreateEpisode(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	ctx, cancel := context.WithCancel(context.Background())
+	ch, cancelSub, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	cancel()
+	drainUntilClosed(t, ch)
+	cancelSub()
+	require.NoError(t, buffer.CreateEpisode(key))
+}
+
+func TestBuffer_SubscribeCanceledWithoutCreateReclaimsEpisode(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancelSub, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	cancelSub()
+	// The subscriber goroutine removes itself from the episode before closing
+	// the output channel, so cleanup is complete once the channel is closed.
+	drainUntilClosed(t, ch)
+
+	_, err = buffer.GetParts(key)
+	require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound)
+	require.Equal(t, 0, buffer.EpisodeCount())
+}
+
+func TestBuffer_CloseClosesPendingSubscriptionAndRejectsOperations(t *testing.T) {
+	t.Parallel()
+
+	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
+	key := testEpisodeKey()
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ch, cancel, err := buffer.SubscribeToEpisode(ctx, key)
+	require.NoError(t, err)
+	defer cancel()
+	buffer.Close()
+	assertChannelClosed(t, ch)
+	require.ErrorIs(t, buffer.CreateEpisode(key), messagepartbuffer.ErrMessagePartBufferClosed)
+}
+
+func testEpisodeKey() messagepartbuffer.Key {
+	return messagepartbuffer.Key{ChatID: uuid.New(), HistoryVersion: 1, GenerationAttempt: 1}
+}
+
+func receivePart(t *testing.T, ch <-chan messagepartbuffer.Part) messagepartbuffer.Part {
+	t.Helper()
+	select {
+	case part, ok := <-ch:
+		require.True(t, ok)
+		return part
+	case <-time.After(testutil.WaitLong):
+		t.Fatal("timed out waiting for buffered part")
+		return messagepartbuffer.Part{}
+	}
+}
+
+func assertChannelClosed[T any](t *testing.T, ch <-chan T) {
+	t.Helper()
+	select {
+	case _, ok := <-ch:
+		require.False(t, ok)
+	case <-time.After(testutil.WaitLong):
+		t.Fatal("timed out waiting for channel close")
+	}
+}
+
+func drainUntilClosed[T any](t *testing.T, ch <-chan T) {
+	t.Helper()
+	for {
+		select {
+		case _, ok := <-ch:
+			if !ok {
+				return
+			}
+		case <-time.After(testutil.WaitLong):
+			t.Fatal("timed out waiting for channel close")
+		}
+	}
+}
+
+func serializedPartBytes(t *testing.T, part messagepartbuffer.Part) int64 {
+	t.Helper()
+	data, err := json.Marshal(struct {
+		Seq  int64                    `json:"seq"`
+		Role codersdk.ChatMessageRole `json:"role"`
+		Part codersdk.ChatMessagePart `json:"part"`
+	}{
+		Seq:  part.Seq,
+		Role: part.Role,
+		Part: part.MessagePart,
+	})
+	require.NoError(t, err)
+	return int64(len(data))
+}