orca.ai/pkg/session/manager.go

package session

import (
	"fmt"
	"sync"
	"time"

	"github.com/orca/orca/pkg/bus"
)

// Manager provides high-level session lifecycle operations.
//
// It wraps a Store with caching, context window management, and
// event publishing on the message bus.
type Manager struct {
	store  Store
	bus    bus.MessageBus
	cache  map[string]*Session
	mu     sync.RWMutex
}

// NewManager creates a new session Manager with the given store and optional message bus.
func NewManager(store Store, mb bus.MessageBus) *Manager {
	return &Manager{
		store: store,
		bus:   mb,
		cache: make(map[string]*Session),
	}
}

// CreateSession creates a new session with the given ID and optional metadata.
func (m *Manager) CreateSession(id string, metadata map[string]string) (*Session, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, exists := m.cache[id]; exists {
		return nil, fmt.Errorf("session %q already exists", id)
	}

	now := time.Now()
	session := &Session{
		ID:        id,
		Status:    SessionActive,
		Messages:  make([]SessionMessage, 0),
		CreatedAt: now,
		UpdatedAt: now,
		Metadata:  metadata,
	}

	m.cache[id] = session

	// Publish session created event
	if m.bus != nil {
		m.bus.Publish("session.created", bus.Message{
			ID:      "session-" + id,
			Type:    bus.MsgTypeSystem,
			From:    "session.manager",
			Content: map[string]interface{}{"session_id": id},
		})
	}

	return session, nil
}

// GetSession retrieves a session by ID, checking the cache and then the store.
func (m *Manager) GetSession(id string) (*Session, error) {
	m.mu.RLock()
	session, ok := m.cache[id]
	m.mu.RUnlock()

	if ok {
		return session, nil
	}

	// Try to load from store
	messages, err := m.store.Load(id)
	if err != nil {
		return nil, fmt.Errorf("failed to load session %q: %w", id, err)
	}

	// Check if we can determine created/updated timestamps from messages
	var createdAt, updatedAt time.Time
	if len(messages) > 0 {
		createdAt = messages[0].Timestamp
		updatedAt = messages[len(messages)-1].Timestamp
	}
	if createdAt.IsZero() {
		createdAt = time.Now()
	}
	if updatedAt.IsZero() {
		updatedAt = time.Now()
	}

	session = &Session{
		ID:        id,
		Status:    SessionActive,
		Messages:  messages,
		CreatedAt: createdAt,
		UpdatedAt: updatedAt,
	}

	m.mu.Lock()
	m.cache[id] = session
	m.mu.Unlock()

	return session, nil
}

// AddMessage appends a message to a session and persists it.
func (m *Manager) AddMessage(sessionID string, role MessageRole, content string, metadata map[string]string) (*SessionMessage, error) {
	msg := SessionMessage{
		Role:      role,
		Content:   content,
		Timestamp: time.Now(),
		Metadata:  metadata,
	}

	if err := m.store.Save(sessionID, msg); err != nil {
		return nil, fmt.Errorf("failed to save message to session %q: %w", sessionID, err)
	}

	// Upsert cache
	m.mu.Lock()
	if session, ok := m.cache[sessionID]; ok {
		session.Messages = append(session.Messages, msg)
		session.UpdatedAt = msg.Timestamp
	} else {
		m.cache[sessionID] = &Session{
			ID:        sessionID,
			Status:    SessionActive,
			Messages:  []SessionMessage{msg},
			CreatedAt: msg.Timestamp,
			UpdatedAt: msg.Timestamp,
		}
	}
	m.mu.Unlock()

	return &msg, nil
}

// GetContext returns the most recent N messages in a session.
// If windowSize <= 0 or >= total messages, all messages are returned.
func (m *Manager) GetContext(sessionID string, windowSize int) ([]SessionMessage, error) {
	session, err := m.GetSession(sessionID)
	if err != nil {
		return nil, err
	}

	messages := session.Messages
	if windowSize > 0 && windowSize < len(messages) {
		return messages[len(messages)-windowSize:], nil
	}
	return messages, nil
}

// ArchiveSession archives a session, making it read-only.
func (m *Manager) ArchiveSession(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if session, ok := m.cache[id]; ok {
		session.Status = SessionArchived
	}

	if err := m.store.Archive(id); err != nil {
		return err
	}

	// Publish event
	if m.bus != nil {
		m.bus.Publish("session.archived", bus.Message{
			ID:      "session-" + id,
			Type:    bus.MsgTypeSystem,
			From:    "session.manager",
			Content: map[string]interface{}{"session_id": id},
		})
	}

	return nil
}

// DeleteSession permanently removes a session.
func (m *Manager) DeleteSession(id string) error {
	m.mu.Lock()
	delete(m.cache, id)
	m.mu.Unlock()
	return m.store.Delete(id)
}

// ListSessions returns all known session IDs.
func (m *Manager) ListSessions() ([]string, error) {
	return m.store.List()
}

// Store returns the underlying Store.
func (m *Manager) Store() Store {
	return m.store
}