orca.ai/pkg/session/jsonl.go

package session

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// JSONLStore implements the Store interface using JSONL files.
//
// Each session is stored in a separate file named {session_id}.jsonl
// under the configured storage directory. Every line in the file is a
// JSON-encoded SessionMessage. New messages are appended in O(1) time.
type JSONLStore struct {
	storageDir string
	mu         sync.RWMutex
}

// NewJSONLStore creates a new JSONLStore with the given storage directory.
// The directory is created if it does not exist.
func NewJSONLStore(storageDir string) (*JSONLStore, error) {
	if err := os.MkdirAll(storageDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create session storage directory %q: %w", storageDir, err)
	}
	return &JSONLStore{storageDir: storageDir}, nil
}

// path returns the full file path for the given session ID.
func (s *JSONLStore) path(sessionID string) string {
	return filepath.Join(s.storageDir, sessionID+".jsonl")
}

// archivePath returns the archive file path for the given session ID.
func (s *JSONLStore) archivePath(sessionID string) string {
	return filepath.Join(s.storageDir, sessionID+".jsonl.archived")
}

// Save appends a message to a session's JSONL file.
// If the file does not exist, it is created.
// This is an O(1) append operation.
func (s *JSONLStore) Save(sessionID string, msg SessionMessage) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	f, err := os.OpenFile(s.path(sessionID), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return fmt.Errorf("failed to open session file for %q: %w", sessionID, err)
	}
	defer f.Close()

	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal session message: %w", err)
	}

	if _, err := f.Write(append(data, '\n')); err != nil {
		return fmt.Errorf("failed to write session message: %w", err)
	}

	return nil
}

// Load retrieves all messages for a session in chronological order.
// Returns an error if the session file does not exist.
func (s *JSONLStore) Load(sessionID string) ([]SessionMessage, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	data, err := os.ReadFile(s.path(sessionID))
	if err != nil {
		if os.IsNotExist(err) {
			// Check archive
			archiveData, archiveErr := os.ReadFile(s.archivePath(sessionID))
			if archiveErr != nil {
				return nil, fmt.Errorf("session %q not found", sessionID)
			}
			data = archiveData
		} else {
			return nil, fmt.Errorf("failed to read session file for %q: %w", sessionID, err)
		}
	}

	return parseJSONL(data)
}

// parseJSONL parses a JSONL byte slice into a slice of SessionMessage.
func parseJSONL(data []byte) ([]SessionMessage, error) {
	var messages []SessionMessage
	trimmed := strings.TrimSpace(string(data))
	if trimmed == "" {
		return messages, nil
	}

	lines := strings.Split(trimmed, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		var msg SessionMessage
		if err := json.Unmarshal([]byte(line), &msg); err != nil {
			return nil, fmt.Errorf("failed to unmarshal session message: %w", err)
		}
		messages = append(messages, msg)
	}
	return messages, nil
}

// List returns all session IDs by scanning the storage directory.
func (s *JSONLStore) List() ([]string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	entries, err := os.ReadDir(s.storageDir)
	if err != nil {
		return nil, fmt.Errorf("failed to read storage directory %q: %w", s.storageDir, err)
	}

	var sessions []string
	for _, entry := range entries {
		name := entry.Name()
		if strings.HasSuffix(name, ".jsonl") && !strings.HasSuffix(name, ".archived") {
			sessions = append(sessions, strings.TrimSuffix(name, ".jsonl"))
		}
	}
	return sessions, nil
}

// Exists checks whether a session file exists (active or archived).
func (s *JSONLStore) Exists(sessionID string) (bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if _, err := os.Stat(s.path(sessionID)); err == nil {
		return true, nil
	} else if !os.IsNotExist(err) {
		return false, fmt.Errorf("failed to check session %q: %w", sessionID, err)
	}

	// Check archive
	if _, err := os.Stat(s.archivePath(sessionID)); err == nil {
		return true, nil
	} else if !os.IsNotExist(err) {
		return false, fmt.Errorf("failed to check archived session %q: %w", sessionID, err)
	}

	return false, nil
}

// Archive moves a session file to the archived state by renaming it.
func (s *JSONLStore) Archive(sessionID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if err := os.Rename(s.path(sessionID), s.archivePath(sessionID)); err != nil {
		if os.IsNotExist(err) {
			return fmt.Errorf("session %q not found", sessionID)
		}
		return fmt.Errorf("failed to archive session %q: %w", sessionID, err)
	}
	return nil
}

// Delete permanently removes a session file and its archive.
func (s *JSONLStore) Delete(sessionID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	var lastErr error

	// Remove active file
	if err := os.Remove(s.path(sessionID)); err != nil && !os.IsNotExist(err) {
		lastErr = fmt.Errorf("failed to delete session %q: %w", sessionID, err)
	}

	// Also remove archived file if it exists
	if err := os.Remove(s.archivePath(sessionID)); err != nil && !os.IsNotExist(err) {
		lastErr = fmt.Errorf("failed to delete archived session %q: %w", sessionID, err)
	}

	return lastErr
}

// StorageDir returns the storage directory path.
func (s *JSONLStore) StorageDir() string {
	return s.storageDir
}