chore: Replace OnSnapshot callback with Emitter interface

lib/httpapi/events.go

@@ -64,7 +64,7 @@ type EventEmitter struct {
 	agentType mf.AgentType
 	chans map[int]chan Event
 	chanIdx int
-	subscriptionBufSize int
+	subscriptionBufSize uint
 	screen string
 }
 
@@ -81,20 +81,37 @@ func convertStatus(status st.ConversationStatus) AgentStatus {
 	}
 }
 
-// subscriptionBufSize is the size of the buffer for each subscription.
-// Once the buffer is full, the channel will be closed.
-// Listeners must actively drain the channel, so it's important to
-// set this to a value that is large enough to handle the expected
-// number of events.
-func NewEventEmitter(subscriptionBufSize int) *EventEmitter {
-	return &EventEmitter{
- mu: sync.Mutex{},
+const defaultSubscriptionBufSize uint = 1024
+
+type EventEmitterOption func(*EventEmitter)
+
+func WithSubscriptionBufSize(size uint) EventEmitterOption {
+	return func(e *EventEmitter) {
+ if size == 0 {
+ e.subscriptionBufSize = defaultSubscriptionBufSize
+ } else {
+ e.subscriptionBufSize = size
+ }
+	}
+}
+
+func WithAgentType(agentType mf.AgentType) EventEmitterOption {
+	return func(e *EventEmitter) {
+ e.agentType = agentType
+	}
+}
+
+func NewEventEmitter(opts ...EventEmitterOption) *EventEmitter {
+	e := &EventEmitter{
  messages: make([]st.ConversationMessage, 0),
  status: AgentStatusRunning,
  chans: make(map[int]chan Event),
- chanIdx: 0,
- subscriptionBufSize: subscriptionBufSize,
+ subscriptionBufSize: defaultSubscriptionBufSize,
+	}
+	for _, opt := range opts {
+ opt(e)
 	}
+	return e
 }
 
 // Assumes the caller holds the lock.
@@ -122,7 +139,7 @@ func (e *EventEmitter) notifyChannels(eventType EventType, payload any) {
 
 // Assumes that only the last message can change or new messages can be added.
 // If a new message is injected between existing messages (identified by Id), the behavior is undefined.
-func (e *EventEmitter) UpdateMessagesAndEmitChanges(newMessages []st.ConversationMessage) {
+func (e *EventEmitter) EmitMessages(newMessages []st.ConversationMessage) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
@@ -137,6 +154,9 @@ func (e *EventEmitter) UpdateMessagesAndEmitChanges(newMessages []st.Conversatio
  newMsg = newMessages[i]
  }
  if oldMsg != newMsg {
+ if i >= len(newMessages) {
+ continue
+ }
  e.notifyChannels(EventTypeMessageUpdate, MessageUpdateBody{
  Id: newMessages[i].Id,
  Role: newMessages[i].Role,
@@ -149,7 +169,7 @@ func (e *EventEmitter) UpdateMessagesAndEmitChanges(newMessages []st.Conversatio
 	e.messages = newMessages
 }
 
-func (e *EventEmitter) UpdateStatusAndEmitChanges(newStatus st.ConversationStatus, agentType mf.AgentType) {
+func (e *EventEmitter) EmitStatus(newStatus st.ConversationStatus) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
@@ -158,12 +178,11 @@ func (e *EventEmitter) UpdateStatusAndEmitChanges(newStatus st.ConversationStatu
  return
 	}
 
-	e.notifyChannels(EventTypeStatusChange, StatusChangeBody{Status: newAgentStatus, AgentType: agentType})
+	e.notifyChannels(EventTypeStatusChange, StatusChangeBody{Status: newAgentStatus, AgentType: e.agentType})
 	e.status = newAgentStatus
-	e.agentType = agentType
 }
 
-func (e *EventEmitter) UpdateScreenAndEmitChanges(newScreen string) {
+func (e *EventEmitter) EmitScreen(newScreen string) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 

lib/httpapi/events_test.go

@@ -5,14 +5,13 @@ import (
 	"testing"
 	"time"
 
-	mf "github.com/coder/agentapi/lib/msgfmt"
 	st "github.com/coder/agentapi/lib/screentracker"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestEventEmitter(t *testing.T) {
 	t.Run("single-subscription", func(t *testing.T) {
- emitter := NewEventEmitter(10)
+ emitter := NewEventEmitter(WithSubscriptionBufSize(10))
  _, ch, stateEvents := emitter.Subscribe()
  assert.Empty(t, ch)
  assert.Equal(t, []Event{
@@ -27,7 +26,7 @@ func TestEventEmitter(t *testing.T) {
  }, stateEvents)
 
  now := time.Now()
- emitter.UpdateMessagesAndEmitChanges([]st.ConversationMessage{
+ emitter.EmitMessages([]st.ConversationMessage{
  {Id: 1, Message: "Hello, world!", Role: st.ConversationRoleUser, Time: now},
  })
  newEvent := <-ch
@@ -36,7 +35,7 @@ func TestEventEmitter(t *testing.T) {
  Payload: MessageUpdateBody{Id: 1, Message: "Hello, world!", Role: st.ConversationRoleUser, Time: now},
  }, newEvent)
 
- emitter.UpdateMessagesAndEmitChanges([]st.ConversationMessage{
+ emitter.EmitMessages([]st.ConversationMessage{
  {Id: 1, Message: "Hello, world! (updated)", Role: st.ConversationRoleUser, Time: now},
  {Id: 2, Message: "What's up?", Role: st.ConversationRoleAgent, Time: now},
  })
@@ -52,24 +51,24 @@ func TestEventEmitter(t *testing.T) {
  Payload: MessageUpdateBody{Id: 2, Message: "What's up?", Role: st.ConversationRoleAgent, Time: now},
  }, newEvent)
 
- emitter.UpdateStatusAndEmitChanges(st.ConversationStatusStable, mf.AgentTypeAider)
+ emitter.EmitStatus(st.ConversationStatusStable)
  newEvent = <-ch
  assert.Equal(t, Event{
  Type: EventTypeStatusChange,
- Payload: StatusChangeBody{Status: AgentStatusStable, AgentType: mf.AgentTypeAider},
+ Payload: StatusChangeBody{Status: AgentStatusStable, AgentType: ""},
  }, newEvent)
 	})
 
 	t.Run("multiple-subscriptions", func(t *testing.T) {
- emitter := NewEventEmitter(10)
+ emitter := NewEventEmitter(WithSubscriptionBufSize(10))
  channels := make([]<-chan Event, 0, 10)
  for i := 0; i < 10; i++ {
  _, ch, _ := emitter.Subscribe()
  channels = append(channels, ch)
  }
  now := time.Now()
 
- emitter.UpdateMessagesAndEmitChanges([]st.ConversationMessage{
+ emitter.EmitMessages([]st.ConversationMessage{
  {Id: 1, Message: "Hello, world!", Role: st.ConversationRoleUser, Time: now},
  })
  for _, ch := range channels {
@@ -82,10 +81,10 @@ func TestEventEmitter(t *testing.T) {
 	})
 
 	t.Run("close-channel", func(t *testing.T) {
- emitter := NewEventEmitter(1)
+ emitter := NewEventEmitter(WithSubscriptionBufSize(1))
  _, ch, _ := emitter.Subscribe()
  for i := range 5 {
- emitter.UpdateMessagesAndEmitChanges([]st.ConversationMessage{
+ emitter.EmitMessages([]st.ConversationMessage{
  {Id: i, Message: fmt.Sprintf("Hello, world! %d", i), Role: st.ConversationRoleUser, Time: time.Now()},
  })
  }

lib/httpapi/server.go

@@ -244,7 +244,7 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
  return mf.FormatToolCall(config.AgentType, message)
 	}
 
-	emitter := NewEventEmitter(1024)
+	emitter := NewEventEmitter(WithAgentType(config.AgentType))
 
 	// Format initial prompt into message parts if provided
 	var initialPrompt []st.MessagePart
@@ -262,16 +262,8 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
  ReadyForInitialPrompt: isAgentReadyForInitialPrompt,
  FormatToolCall: formatToolCall,
  InitialPrompt: initialPrompt,
- // OnSnapshot uses a callback rather than passing the emitter directly
- // to keep the screentracker package decoupled from httpapi concerns.
- // This preserves clean package boundaries and avoids import cycles.
- OnSnapshot: func(status st.ConversationStatus, messages []st.ConversationMessage, screen string) {
- emitter.UpdateStatusAndEmitChanges(status, config.AgentType)
- emitter.UpdateMessagesAndEmitChanges(messages)
- emitter.UpdateScreenAndEmitChanges(screen)
- },
- Logger: logger,
-	})
+ Logger: logger,
+	}, emitter)
 
 	// Create temporary directory for uploads
 	tempDir, err := os.MkdirTemp("", "agentapi-uploads-")

lib/screentracker/conversation.go

@@ -65,6 +65,13 @@ type Conversation interface {
 	Text() string
 }
 
+// Emitter receives conversation state updates.
+type Emitter interface {
+	EmitMessages([]ConversationMessage)
+	EmitStatus(ConversationStatus)
+	EmitScreen(string)
+}
+
 type ConversationMessage struct {
 	Id int
 	Message string

lib/screentracker/pty_conversation.go

@@ -68,9 +68,7 @@ type PTYConversationConfig struct {
 	FormatToolCall func(message string) (string, []string)
 	// InitialPrompt is the initial prompt to send to the agent once ready
 	InitialPrompt []MessagePart
-	// OnSnapshot is called after each snapshot with current status, messages, and screen content
-	OnSnapshot func(status ConversationStatus, messages []ConversationMessage, screen string)
-	Logger *slog.Logger
+	Logger *slog.Logger
 }
 
 func (cfg PTYConversationConfig) getStableSnapshotsThreshold() int {
@@ -86,7 +84,8 @@ func (cfg PTYConversationConfig) getStableSnapshotsThreshold() int {
 // PTYConversation is a conversation that uses a pseudo-terminal (PTY) for communication.
 // It uses a combination of polling and diffs to detect changes in the screen.
 type PTYConversation struct {
-	cfg PTYConversationConfig
+	cfg PTYConversationConfig
+	emitter Emitter
 	// How many stable snapshots are required to consider the screen stable
 	stableSnapshotsThreshold int
 	snapshotBuffer *RingBuffer[screenSnapshot]
@@ -115,13 +114,23 @@ type PTYConversation struct {
 
 var _ Conversation = &PTYConversation{}
 
-func NewPTY(ctx context.Context, cfg PTYConversationConfig) *PTYConversation {
+type noopEmitter struct{}
+
+func (noopEmitter) EmitMessages([]ConversationMessage) {}
+func (noopEmitter) EmitStatus(ConversationStatus) {}
+func (noopEmitter) EmitScreen(string) {}
+
+func NewPTY(ctx context.Context, cfg PTYConversationConfig, emitter Emitter) *PTYConversation {
 	if cfg.Clock == nil {
  cfg.Clock = quartz.NewReal()
 	}
+	if emitter == nil {
+ emitter = noopEmitter{}
+	}
 	threshold := cfg.getStableSnapshotsThreshold()
 	c := &PTYConversation{
  cfg: cfg,
+ emitter: emitter,
  stableSnapshotsThreshold: threshold,
  snapshotBuffer: NewRingBuffer[screenSnapshot](threshold),
  messages: []ConversationMessage{
@@ -139,9 +148,6 @@ func NewPTY(ctx context.Context, cfg PTYConversationConfig) *PTYConversation {
 	if len(cfg.InitialPrompt) > 0 {
  c.outboundQueue <- outboundMessage{parts: cfg.InitialPrompt, errCh: nil}
 	}
-	if c.cfg.OnSnapshot == nil {
- c.cfg.OnSnapshot = func(ConversationStatus, []ConversationMessage, string) {}
-	}
 	if c.cfg.ReadyForInitialPrompt == nil {
  c.cfg.ReadyForInitialPrompt = func(string) bool { return true }
 	}
@@ -173,7 +179,9 @@ func (c *PTYConversation) Start(ctx context.Context) {
  }
  c.lock.Unlock()
 
- c.cfg.OnSnapshot(status, messages, screen)
+ c.emitter.EmitStatus(status)
+ c.emitter.EmitMessages(messages)
+ c.emitter.EmitScreen(screen)
  return nil
 	}, "snapshot")