session.go

package flowengine

import (
	"sync"
	"time"
)

// ===========================================================================
// Session - Session Context Management
// ===========================================================================

// Session represents a session with context chaining.
type Session struct {
	// id is the session root identifier.
	id string

	// rootResponseID is the LLM response ID for the root node (cache prefix).
	rootResponseID string

	// nodeResponses maps step names to LLM response IDs.
	nodeResponses map[string]string

	// chain records execution order for cleanup.
	chain []string

	// messages accumulates session history.
	messages []Message

	// lastAccessed records the last access time.
	lastAccessed time.Time

	mu sync.RWMutex
}

// Message represents a session message.
type Message struct {
	Role    string `json:"role"` // system, user, assistant
	Content string `json:"content"`
}

// NewSession creates a new session with the given ID.
func NewSession(id string) *Session {
	return &Session{
		id:            id,
		nodeResponses: make(map[string]string),
		chain:         make([]string, 0),
		messages:      make([]Message, 0),
		lastAccessed:  time.Now(),
	}
}

// touch updates the last access time.
func (s *Session) touch() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.touchLocked()
}

func (s *Session) touchLocked() {
	s.lastAccessed = time.Now()
}

// LastAccessed returns the session's last access time.
func (s *Session) LastAccessed() time.Time {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.lastAccessed
}

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

// SetRootResponseID sets the root LLM response ID.
func (s *Session) SetRootResponseID(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.rootResponseID = id
	s.touchLocked()
}

// RootResponseID returns the root LLM response ID.
func (s *Session) RootResponseID() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.rootResponseID
}

// RecordNode records step execution with its LLM response ID.
func (s *Session) RecordNode(stepName, responseID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.nodeResponses[stepName] = responseID
	s.chain = append(s.chain, stepName)
	s.touchLocked()
}

// GetResponseID returns the LLM response ID for a step.
func (s *Session) GetResponseID(stepName string) (string, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	id, ok := s.nodeResponses[stepName]
	return id, ok
}

// LastResponseID returns the most recent LLM response ID.
func (s *Session) LastResponseID() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if len(s.chain) == 0 {
		return s.rootResponseID
	}
	lastStep := s.chain[len(s.chain)-1]
	return s.nodeResponses[lastStep]
}

// AppendMessage adds a message to session history.
func (s *Session) AppendMessage(role, content string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.messages = append(s.messages, Message{Role: role, Content: content})
	s.touchLocked()
}

// Messages returns all session messages.
func (s *Session) Messages() []Message {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]Message, len(s.messages))
	copy(result, s.messages)
	return result
}

// Chain returns the execution chain.
func (s *Session) Chain() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	result := make([]string, len(s.chain))
	copy(result, s.chain)
	return result
}

// Clear resets the session state.
func (s *Session) Clear() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.rootResponseID = ""
	s.nodeResponses = make(map[string]string)
	s.chain = make([]string, 0)
	s.messages = make([]Message, 0)
	s.touchLocked()
}

// ===========================================================================
// SessionManager - Manages Multiple Sessions with TTL and Size Limits
// ===========================================================================

// DefaultSessionTTL is the default session TTL (30 minutes).
const DefaultSessionTTL = 30 * time.Minute

// DefaultMaxSessions is the default maximum session count.
const DefaultMaxSessions = 1000

// SessionManager manages sessions.
type SessionManager struct {
	sessions  map[string]*Session
	ttl       time.Duration
	maxSize   int
	mu        sync.RWMutex
	stopClean chan struct{}
	cleanOnce sync.Once
}

// SessionManagerOption configures SessionManager.
type SessionManagerOption func(*SessionManager)

// WithTTL sets the session TTL.
func WithTTL(ttl time.Duration) SessionManagerOption {
	return func(m *SessionManager) {
		m.ttl = ttl
	}
}

// WithMaxSessions sets the maximum session count.
func WithMaxSessions(max int) SessionManagerOption {
	return func(m *SessionManager) {
		m.maxSize = max
	}
}

// NewSessionManager creates a new session manager.
func NewSessionManager(opts ...SessionManagerOption) *SessionManager {
	m := &SessionManager{
		sessions:  make(map[string]*Session),
		ttl:       DefaultSessionTTL,
		maxSize:   DefaultMaxSessions,
		stopClean: make(chan struct{}),
	}
	for _, opt := range opts {
		opt(m)
	}
	// Start background cleanup goroutine
	go m.cleanupLoop()
	return m
}

// cleanupLoop periodically removes expired sessions.
func (m *SessionManager) cleanupLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-m.stopClean:
			return
		case <-ticker.C:
			m.cleanup()
		}
	}
}

// cleanup removes expired sessions and enforces size limits.
func (m *SessionManager) cleanup() {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	var toDelete []string

	// Find expired sessions
	for id, sess := range m.sessions {
		if now.Sub(sess.LastAccessed()) > m.ttl {
			toDelete = append(toDelete, id)
		}
	}

	// Delete expired sessions
	for _, id := range toDelete {
		delete(m.sessions, id)
	}

	// Enforce max size by removing oldest sessions
	if len(m.sessions) > m.maxSize {
		excess := len(m.sessions) - m.maxSize
		type sessionAge struct {
			id   string
			time time.Time
		}
		ages := make([]sessionAge, 0, len(m.sessions))
		for id, sess := range m.sessions {
			ages = append(ages, sessionAge{id: id, time: sess.LastAccessed()})
		}
		// Sort oldest first (simple selection for small overflow)
		for i := 0; i < excess; i++ {
			oldest := i
			for j := i + 1; j < len(ages); j++ {
				if ages[j].time.Before(ages[oldest].time) {
					oldest = j
				}
			}
			ages[i], ages[oldest] = ages[oldest], ages[i]
			delete(m.sessions, ages[i].id)
		}
	}
}

// Stop stops the cleanup goroutine.
func (m *SessionManager) Stop() {
	m.cleanOnce.Do(func() {
		close(m.stopClean)
	})
}

// Get returns or creates a session.
func (m *SessionManager) Get(id string) *Session {
	m.mu.Lock()
	defer m.mu.Unlock()
	if s, ok := m.sessions[id]; ok {
		s.touch()
		return s
	}
	s := NewSession(id)
	m.sessions[id] = s
	return s
}

// GetIfExists returns a session if it exists, nil otherwise.
func (m *SessionManager) GetIfExists(id string) *Session {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if s, ok := m.sessions[id]; ok {
		return s
	}
	return nil
}

// Delete removes a session and all its contexts.
func (m *SessionManager) Delete(id string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.sessions, id)
}

// Exists checks if a session exists.
func (m *SessionManager) Exists(id string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	_, ok := m.sessions[id]
	return ok
}

// Count returns the number of active sessions.
func (m *SessionManager) Count() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.sessions)
}

// Clear removes all sessions.
func (m *SessionManager) Clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessions = make(map[string]*Session)
}

// TTL returns the current TTL setting.
func (m *SessionManager) TTL() time.Duration {
	return m.ttl
}

// MaxSize returns the current max size setting.
func (m *SessionManager) MaxSize() int {
	return m.maxSize
}