zero/internal/daemon/session.go

package daemon

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"
)

// Session/lease/roster model. Mirrors reference-daemon-code-agent-js/
// session-manager.js (sessionID -> worker + metrics), lease.js (one active
// request per worker) and roster.js (the live session registry). Because zero's
// `exec` worker is one-shot, each session is served by exactly one worker via
// Pool.Run, so "one active request per worker" holds inherently and the pool's
// bounded slots provide the queue-when-full behavior.

// SessionState is a session's lifecycle phase.
type SessionState string

const (
	// SessionQueued: created, waiting for a free pool slot.
	SessionQueued SessionState = "queued"
	// SessionRunning: a worker is leased and producing output.
	SessionRunning SessionState = "running"
	// SessionDone: the worker finished cleanly.
	SessionDone SessionState = "done"
	// SessionFailed: the worker failed permanently / the run errored.
	SessionFailed SessionState = "failed"
)

// ErrSessionExists is returned when starting a session ID already in use.
var ErrSessionExists = errors.New("daemon: session already exists")

// ErrSessionNotFound is returned by Get/Attach for an unknown session ID.
var ErrSessionNotFound = errors.New("daemon: session not found")

const defaultSessionBuffer = 1024

// Session is one routed agent run. It buffers recent output (so a late `attach`
// sees history) and fans live lines out to subscribers. It implements the pool's
// Sink (Line) and the lease "started" hook (Started).
type Session struct {
	id  string
	cwd string

	mu          sync.Mutex
	state       SessionState
	buffer      []string
	maxBuffer   int
	lines       int
	finished    bool
	err         error
	exitCode    int
	subscribers map[int]*subscriber
	nextSub     int
	done        chan struct{}
}

// subscriber is one live attach: a bounded channel plus a count of live lines
// dropped because that channel was full. The drop count is flushed to the consumer
// as a gap notice (see Line) the moment the channel drains, so a lagging client
// learns its stream is incomplete instead of silently losing output (D8).
type subscriber struct {
	ch      chan string
	dropped int
}

// gapNotice is a stream-json event line emitted to a subscriber that fell behind,
// reporting how many live lines it missed. It uses a daemon-namespaced type so a
// stream-json consumer that switches on "type" ignores it rather than mistaking it
// for a terminal event; the full output remains available via a fresh attach
// (which replays the buffer).
func gapNotice(dropped int) string {
	return fmt.Sprintf(`{"type":"daemon_gap","dropped":%d}`, dropped)
}

func newSession(id, cwd string, maxBuffer int) *Session {
	if maxBuffer <= 0 {
		maxBuffer = defaultSessionBuffer
	}
	return &Session{
		id:          id,
		cwd:         cwd,
		state:       SessionQueued,
		maxBuffer:   maxBuffer,
		subscribers: map[int]*subscriber{},
		done:        make(chan struct{}),
	}
}

// ID returns the session ID.
func (s *Session) ID() string { return s.id }

// Done is closed when the session finishes (success or failure).
func (s *Session) Done() <-chan struct{} { return s.done }

// Err returns the terminal error (nil on clean completion); valid after Done.
func (s *Session) Err() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.err
}

// State returns the current lifecycle phase.
func (s *Session) State() SessionState {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.state
}

// Started marks the session running once the pool leases a worker (lease hook).
func (s *Session) Started() {
	s.mu.Lock()
	if s.state == SessionQueued {
		s.state = SessionRunning
	}
	s.mu.Unlock()
}

// Line records and broadcasts one stream-json output line.
func (s *Session) Line(line string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.lines++
	s.buffer = append(s.buffer, line)
	if len(s.buffer) > s.maxBuffer {
		// Drop the oldest line; the ring keeps only the most recent maxBuffer.
		s.buffer = s.buffer[len(s.buffer)-s.maxBuffer:]
	}
	// Hold s.mu across the (non-blocking) sends. The sends never block — every case
	// has a default — so the lock is held only briefly, and holding it stops
	// cancel()/finish() from closing a subscriber channel mid-send, which would
	// panic with "send on closed channel". It also makes each subscriber's dropped
	// counter mutated under the lock.
	for _, sub := range s.subscribers {
		// Non-blocking: a slow/stalled subscriber must not stall the worker pump.
		if sub.dropped > 0 {
			// The consumer fell behind earlier. Before delivering more output, flush a
			// gap notice so it learns its stream is incomplete (D8). If the channel is
			// still full, keep counting and skip this line too.
			select {
			case sub.ch <- gapNotice(sub.dropped):
				sub.dropped = 0
			default:
				sub.dropped++
				continue
			}
		}
		select {
		case sub.ch <- line:
		default:
			// Channel full: count the drop; the gap notice above will report it once
			// the consumer catches up. Full output stays available via a fresh attach.
			sub.dropped++
		}
	}
}

func (s *Session) finish(exitCode int, err error) {
	s.mu.Lock()
	s.finished = true
	s.exitCode = exitCode
	s.err = err
	if err != nil {
		s.state = SessionFailed
	} else {
		s.state = SessionDone
	}
	subs := s.subscribers
	s.subscribers = map[int]*subscriber{}
	s.mu.Unlock()
	for _, sub := range subs {
		close(sub.ch)
	}
	close(s.done)
}

// Subscribe returns the buffered history plus a channel of subsequent live lines
// and a cancel func. If the session has already finished, the live channel is
// returned already closed. Mirrors attach.js.
func (s *Session) Subscribe() (buffered []string, live <-chan string, cancel func()) {
	s.mu.Lock()
	defer s.mu.Unlock()
	buffered = append([]string(nil), s.buffer...)
	if s.finished {
		ch := make(chan string)
		close(ch)
		return buffered, ch, func() {}
	}
	id := s.nextSub
	s.nextSub++
	sub := &subscriber{ch: make(chan string, 256)}
	s.subscribers[id] = sub
	cancel = func() {
		s.mu.Lock()
		if c, ok := s.subscribers[id]; ok {
			delete(s.subscribers, id)
			close(c.ch)
		}
		s.mu.Unlock()
	}
	return buffered, sub.ch, cancel
}

func (s *Session) status() SessionStatus {
	s.mu.Lock()
	defer s.mu.Unlock()
	return SessionStatus{ID: s.id, State: string(s.state), Lines: s.lines}
}

func (s *Session) isFinished() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.finished
}

// defaultMaxSessions bounds the retained session registry so a long-running
// daemon does not accumulate finished sessions without limit.
const defaultMaxSessions = 256

// SessionManager owns the live session registry and routes each session to a
// pool worker. Mirrors session-manager.js + roster.js. Finished sessions are
// retained (so a late `attach` sees history) up to MaxSessions, past which the
// oldest FINISHED ones are evicted; running/queued sessions are never evicted.
type SessionManager struct {
	pool        *Pool
	maxBuffer   int
	maxSessions int

	mu       sync.Mutex
	sessions map[string]*Session
	order    []string // creation order, for FIFO eviction of finished sessions
}

// SessionManagerOptions configures a SessionManager.
type SessionManagerOptions struct {
	Pool      *Pool
	MaxBuffer int // per-session output ring size; 0 => default
	// MaxSessions caps the retained session registry; 0 => default. The oldest
	// FINISHED sessions are evicted once the cap is exceeded.
	MaxSessions int
}

// NewSessionManager builds a manager over pool.
func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) {
	if opts.Pool == nil {
		return nil, errors.New("daemon: session manager requires a Pool")
	}
	maxSessions := opts.MaxSessions
	if maxSessions <= 0 {
		maxSessions = defaultMaxSessions
	}
	return &SessionManager{
		pool:        opts.Pool,
		maxBuffer:   opts.MaxBuffer,
		maxSessions: maxSessions,
		sessions:    map[string]*Session{},
	}, nil
}

// Start creates a session for spec and dispatches it to the pool. It returns the
// Session immediately (non-blocking); the run proceeds in the background and the
// session transitions queued -> running -> done/failed. A duplicate ID is
// rejected with ErrSessionExists.
func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, error) {
	m.mu.Lock()
	if _, exists := m.sessions[spec.Session]; exists {
		m.mu.Unlock()
		return nil, ErrSessionExists
	}
	sess := newSession(spec.Session, spec.Cwd, m.maxBuffer)
	m.sessions[spec.Session] = sess
	m.order = append(m.order, spec.Session)
	m.pruneLocked()
	m.mu.Unlock()

	go func() {
		code, err := m.pool.Run(ctx, spec, sess)
		sess.finish(code, err)
	}()
	return sess, nil
}

// pruneLocked evicts the oldest FINISHED sessions while the registry exceeds
// maxSessions, so a long-running daemon does not retain finished sessions without
// bound. Running/queued sessions are never evicted (a transient burst can exceed
// the cap until those finish). Caller must hold m.mu.
func (m *SessionManager) pruneLocked() {
	if len(m.sessions) <= m.maxSessions {
		return
	}
	kept := m.order[:0]
	for _, id := range m.order {
		s, ok := m.sessions[id]
		if !ok {
			continue // already removed
		}
		if len(m.sessions) > m.maxSessions && s.isFinished() {
			delete(m.sessions, id)
			continue
		}
		kept = append(kept, id)
	}
	m.order = kept
}

// Get returns a session by ID.
func (m *SessionManager) Get(id string) (*Session, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	s, ok := m.sessions[id]
	return s, ok
}

// Attach returns the buffered history + live stream for an existing session.
func (m *SessionManager) Attach(id string) (buffered []string, live <-chan string, cancel func(), err error) {
	s, ok := m.Get(id)
	if !ok {
		return nil, nil, nil, ErrSessionNotFound
	}
	b, l, c := s.Subscribe()
	return b, l, c, nil
}

// Statuses returns a snapshot of all sessions, sorted by ID for stable output.
func (m *SessionManager) Statuses() []SessionStatus {
	m.mu.Lock()
	sessions := make([]*Session, 0, len(m.sessions))
	for _, s := range m.sessions {
		sessions = append(sessions, s)
	}
	m.mu.Unlock()
	out := make([]SessionStatus, 0, len(sessions))
	for _, s := range sessions {
		out = append(out, s.status())
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out
}