enkly-ai/services/agent/internal/domain/stream.go

// Package domain holds the agent service's own aggregate: one in-flight
// (or just-finished) agent run. Transplanted from
// internal/handlers/chat.go's StreamManager/Stream, with one
// deliberate change: where the monolith's Stream stored pre-rendered
// HTML (because the same process also owned templates), this one
// stores plain text and structured tool calls — rendering now happens
// in the gateway, the only service that imports pkg/renderkit. See
// ARCHITECTURE.md's "Where rendering lives" note.
package domain

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"sync"
	"time"
)

const (
	MaxRunDuration    = 30 * time.Minute // hard cap on one agent turn
	MaxStreamsPerUser = 3                // concurrent in-flight replies per account
	FinishedStreamTTL = 10 * time.Minute // how long a finished stream stays replayable
)

var (
	ErrConversationBusy = errors.New("this conversation is still answering")
	ErrTooManyStreams   = errors.New("too many replies in progress")
)

// ToolCallResult is one completed tool call, as recorded on a Stream.
type ToolCallResult struct {
	Name          string
	ArgumentsJSON string
	Result        string
	IsError       bool
}

// ThinkingRound is one agent round's own standalone reasoning trace, as
// recorded on a Stream — the round-indexed counterpart to ThinkText's
// single running total. See agentrun.Callbacks.OnThinkingRound.
type ThinkingRound struct {
	Round    int
	Thinking string
}

// Stream is the observable state of one agent run. Like the monolith's
// version, it keeps the latest snapshot rather than a log of every
// token event (token text is cumulative anyway), plus an append-only
// list of completed tool calls, keeping memory bounded and letting a
// late attacher (StreamTurn reconnecting) cheaply catch up.
type Stream struct {
	ID     string
	UserID int64
	ConvID int64
	cancel context.CancelFunc

	mu             sync.Mutex
	changed        chan struct{} // closed & replaced on every update (broadcast)
	tokenText      string        // latest cumulative reply
	thinkText      string        // latest cumulative reasoning trace
	tools          []ToolCallResult
	thinkingRounds []ThinkingRound // append-only, one per round that produced reasoning — see AddThinkingRound
	finalText      string
	// finalToolCalls is every tool call made during the run (successful
	// or not), handed to the gateway alongside finalText so it can
	// build deliverable cards the same way
	// internal/handlers/deliverables.go did — see TurnDone in
	// api/proto/agent/v1/agent.proto.
	finalToolCalls []ToolCallResult
	done           bool
}

type StreamSnapshot struct {
	TokenText, ThinkText, FinalText string
	Tools                           []ToolCallResult
	ThinkingRounds                  []ThinkingRound
	FinalToolCalls                  []ToolCallResult
	Done                            bool
}

func (s *Stream) update(fn func()) {
	s.mu.Lock()
	fn()
	close(s.changed)
	s.changed = make(chan struct{})
	s.mu.Unlock()
}

func (s *Stream) Snapshot() (StreamSnapshot, <-chan struct{}) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return StreamSnapshot{
		TokenText: s.tokenText, ThinkText: s.thinkText, FinalText: s.finalText,
		Tools:          s.tools[:len(s.tools):len(s.tools)],                            // append-only; safe to share
		ThinkingRounds: s.thinkingRounds[:len(s.thinkingRounds):len(s.thinkingRounds)], // append-only; safe to share
		FinalToolCalls: s.finalToolCalls,
		Done:           s.done,
	}, s.changed
}

func (s *Stream) IsDone() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.done
}

func (s *Stream) SetToken(text string)      { s.update(func() { s.tokenText = text }) }
func (s *Stream) SetThinking(text string)   { s.update(func() { s.thinkText = text }) }
func (s *Stream) AddTool(tc ToolCallResult) { s.update(func() { s.tools = append(s.tools, tc) }) }

// AddThinkingRound records one round's standalone reasoning trace —
// the round-indexed counterpart to SetThinking's running total. See
// agentrun.Callbacks.OnThinkingRound.
func (s *Stream) AddThinkingRound(round int, thinking string) {
	s.update(func() { s.thinkingRounds = append(s.thinkingRounds, ThinkingRound{Round: round, Thinking: thinking}) })
}
func (s *Stream) Finish(finalText string) {
	s.update(func() { s.finalText = finalText; s.done = true })
}

// FinishWithToolCalls is Finish plus the full tool-call list for this
// run, used by the normal (non-panic, non-empty-history) completion
// path in runTurn.
func (s *Stream) FinishWithToolCalls(finalText string, toolCalls []ToolCallResult) {
	s.update(func() { s.finalText = finalText; s.finalToolCalls = toolCalls; s.done = true })
}

func (s *Stream) Cancel() { s.cancel() }

// Manager tracks in-flight (and just-finished) agent runs — identical
// in behavior to the monolith's StreamManager, including the "second
// concurrent run on the same conversation is refused, not queued"
// invariant (their messages would interleave) and the per-user
// concurrency cap.
type Manager struct {
	mu      sync.Mutex
	streams map[string]*Stream
	byConv  map[int64]*Stream
}

func NewManager() *Manager {
	return &Manager{streams: map[string]*Stream{}, byConv: map[int64]*Stream{}}
}

func (m *Manager) Create(userID, convID int64, cancel context.CancelFunc) (*Stream, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.byConv[convID]; ok && !s.IsDone() {
		return nil, ErrConversationBusy
	}
	active := 0
	for _, s := range m.streams {
		if s.UserID == userID && !s.IsDone() {
			active++
		}
	}
	if active >= MaxStreamsPerUser {
		return nil, ErrTooManyStreams
	}

	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return nil, err
	}
	s := &Stream{ID: hex.EncodeToString(b[:]), UserID: userID, ConvID: convID, cancel: cancel, changed: make(chan struct{})}
	m.streams[s.ID] = s
	m.byConv[convID] = s
	return s, nil
}

func (m *Manager) Get(id string) *Stream {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.streams[id]
}

func (m *Manager) ActiveFor(convID int64) *Stream {
	m.mu.Lock()
	defer m.mu.Unlock()
	if s, ok := m.byConv[convID]; ok && !s.IsDone() {
		return s
	}
	return nil
}

func (m *Manager) CancelConversation(convID int64) {
	if s := m.ActiveFor(convID); s != nil {
		s.Cancel()
	}
}

func (m *Manager) ScheduleRemoval(s *Stream) {
	time.AfterFunc(FinishedStreamTTL, func() {
		m.mu.Lock()
		defer m.mu.Unlock()
		delete(m.streams, s.ID)
		if m.byConv[s.ConvID] == s {
			delete(m.byConv, s.ConvID)
		}
	})
}