codex-agent-manager/internal/runtime/store.go

package runtime

import (
	"database/sql"
	"errors"
	"net/url"
	"os"
	"strings"

	"codex-agent-manager/internal/codexhome"

	_ "modernc.org/sqlite"
)

type Store struct {
	CodexHome string
}

func (s Store) Snapshot() (Snapshot, error) {
	statePath, err := codexhome.ResolveInside(s.CodexHome, "state_5.sqlite")
	if err != nil {
		return Snapshot{}, err
	}
	goalsPath, err := codexhome.ResolveInside(s.CodexHome, "goals_1.sqlite")
	if err != nil {
		return Snapshot{}, err
	}
	stateExists := fileExists(statePath)
	goalsExists := fileExists(goalsPath)
	sources := SourceMap{
		"state": sqliteSource("state", statePath, stateExists),
		"goals": sqliteSource("goals", goalsPath, goalsExists),
	}
	snapshot := Snapshot{
		Threads:    []Thread{},
		SpawnEdges: []SpawnEdge{},
		Goals:      []Goal{},
		Source:     aggregateSource(sources),
		Sources:    sources,
	}
	if stateExists {
		db, err := openReadonlySQLite(statePath)
		if err != nil {
			return Snapshot{}, err
		}
		defer db.Close()
		snapshot.Threads, err = readThreads(db, statePath, snapshot.Sources)
		if err != nil {
			return Snapshot{}, err
		}
		snapshot.SpawnEdges, err = readSpawnEdges(db, statePath, snapshot.Sources)
		if err != nil {
			return Snapshot{}, err
		}
	}
	if goalsExists {
		db, err := openReadonlySQLite(goalsPath)
		if err != nil {
			return Snapshot{}, err
		}
		defer db.Close()
		snapshot.Goals, err = readGoals(db, goalsPath, snapshot.Sources)
		if err != nil {
			return Snapshot{}, err
		}
	}
	return snapshot, nil
}

func sqliteSource(name string, path string, exists bool) SourceEvidence {
	if exists {
		return SourceEvidence{Kind: "sqlite_readonly", Path: path, Confidence: "high"}
	}
	return SourceEvidence{
		Kind:       "sqlite_missing",
		Path:       path,
		Confidence: "low",
		Message:    name + " SQLite file was not found; returning empty data for that source.",
	}
}

func aggregateSource(sources SourceMap) SourceEvidence {
	state := sources["state"]
	goals := sources["goals"]
	if state.Kind == "sqlite_readonly" && goals.Kind == "sqlite_readonly" {
		return SourceEvidence{Kind: "sqlite_readonly", Confidence: "high"}
	}
	if state.Kind == "sqlite_missing" && goals.Kind == "sqlite_missing" {
		return SourceEvidence{
			Kind:       "sqlite_missing",
			Confidence: "low",
			Message:    "Codex SQLite files were not found; returning an empty read-only snapshot.",
		}
	}
	return SourceEvidence{
		Kind:       "sqlite_partial",
		Confidence: "medium",
		Message:    "One Codex SQLite source is missing; available data was read from existing sources only.",
	}
}

func openReadonlySQLite(path string) (*sql.DB, error) {
	uri := url.URL{Scheme: "file", Path: path}
	query := uri.Query()
	query.Set("mode", "ro")
	query.Set("immutable", "1")
	uri.RawQuery = query.Encode()
	return sql.Open("sqlite", uri.String())
}

func readThreads(db *sql.DB, sourcePath string, sources SourceMap) ([]Thread, error) {
	columns, err := tableColumns(db, "threads")
	if err != nil {
		return nil, err
	}
	if len(columns) == 0 {
		sources["threads"] = tableSource("sqlite_missing_table", sourcePath, "threads table was not found.")
		return []Thread{}, nil
	}
	if !columns["id"] {
		sources["threads"] = tableSource("sqlite_schema_drift", sourcePath, "threads table is missing required id column.")
		return []Thread{}, nil
	}
	query := `SELECT ` +
		textColumn(columns, "id") + `, ` +
		firstTextColumn(columns, "agent_role", "role") + `, ` +
		textColumn(columns, "status") + `, ` +
		firstTextColumn(columns, "created_at_ms", "created_at") + `, ` +
		firstTextColumn(columns, "updated_at_ms", "updated_at") + `, ` +
		textColumn(columns, "cwd") + `, ` +
		textColumn(columns, "title") + `, ` +
		textColumn(columns, "agent_nickname") + `, ` +
		textColumn(columns, "agent_role") + `, ` +
		textColumn(columns, "agent_path") + `, ` +
		textColumn(columns, "thread_source") + `, ` +
		textColumn(columns, "preview") +
		` FROM threads ORDER BY ` + orderBy(columns, "created_at_ms", "created_at", "id")
	rows, err := db.Query(query)
	if err != nil {
		if isMissingTable(err) {
			return []Thread{}, nil
		}
		return nil, err
	}
	defer rows.Close()
	var threads []Thread
	for rows.Next() {
		var item Thread
		if err := rows.Scan(
			&item.ID,
			&item.Role,
			&item.Status,
			&item.CreatedAt,
			&item.UpdatedAt,
			&item.CWD,
			&item.Title,
			&item.AgentNickname,
			&item.AgentRole,
			&item.AgentPath,
			&item.ThreadSource,
			&item.Preview,
		); err != nil {
			return nil, err
		}
		item.Source = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
		threads = append(threads, item)
	}
	sources["threads"] = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
	return threads, rows.Err()
}

func readSpawnEdges(db *sql.DB, sourcePath string, sources SourceMap) ([]SpawnEdge, error) {
	columns, err := tableColumns(db, "thread_spawn_edges")
	if err != nil {
		return nil, err
	}
	if len(columns) == 0 {
		sources["thread_spawn_edges"] = tableSource("sqlite_missing_table", sourcePath, "thread_spawn_edges table was not found.")
		return []SpawnEdge{}, nil
	}
	if !(columns["from_thread_id"] && columns["to_thread_id"]) && !(columns["parent_thread_id"] && columns["child_thread_id"]) {
		sources["thread_spawn_edges"] = tableSource("sqlite_schema_drift", sourcePath, "thread_spawn_edges table is missing required endpoint columns.")
		return []SpawnEdge{}, nil
	}
	query := `SELECT ` +
		firstTextColumn(columns, "from_thread_id", "parent_thread_id") + `, ` +
		firstTextColumn(columns, "to_thread_id", "child_thread_id") + `, ` +
		firstTextColumn(columns, "reason", "status") + `, ` +
		textColumn(columns, "created_at") +
		` FROM thread_spawn_edges ORDER BY ` + orderBy(columns, "created_at", "from_thread_id", "parent_thread_id", "to_thread_id", "child_thread_id")
	rows, err := db.Query(query)
	if err != nil {
		if isMissingTable(err) {
			return []SpawnEdge{}, nil
		}
		return nil, err
	}
	defer rows.Close()
	var edges []SpawnEdge
	for rows.Next() {
		var item SpawnEdge
		if err := rows.Scan(&item.FromThreadID, &item.ToThreadID, &item.Reason, &item.CreatedAt); err != nil {
			return nil, err
		}
		item.Source = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
		edges = append(edges, item)
	}
	sources["thread_spawn_edges"] = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
	return edges, rows.Err()
}

func readGoals(db *sql.DB, sourcePath string, sources SourceMap) ([]Goal, error) {
	columns, err := tableColumns(db, "thread_goals")
	if err != nil {
		return nil, err
	}
	if len(columns) == 0 {
		sources["thread_goals"] = tableSource("sqlite_missing_table", sourcePath, "thread_goals table was not found.")
		return []Goal{}, nil
	}
	if !columns["thread_id"] {
		sources["thread_goals"] = tableSource("sqlite_schema_drift", sourcePath, "thread_goals table is missing required thread_id column.")
		return []Goal{}, nil
	}
	query := `SELECT ` +
		textColumn(columns, "thread_id") + `, ` +
		firstTextColumn(columns, "goal", "objective") + `, ` +
		textColumn(columns, "status") + `, ` +
		firstTextColumn(columns, "updated_at_ms", "updated_at") +
		` FROM thread_goals ORDER BY ` + orderBy(columns, "updated_at_ms", "updated_at", "thread_id")
	rows, err := db.Query(query)
	if err != nil {
		if isMissingTable(err) {
			return []Goal{}, nil
		}
		return nil, err
	}
	defer rows.Close()
	var goals []Goal
	for rows.Next() {
		var item Goal
		if err := rows.Scan(&item.ThreadID, &item.Goal, &item.Status, &item.UpdatedAt); err != nil {
			return nil, err
		}
		item.Source = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
		goals = append(goals, item)
	}
	sources["thread_goals"] = SourceEvidence{Kind: "sqlite_table", Path: sourcePath, Confidence: "high"}
	return goals, rows.Err()
}

func tableColumns(db *sql.DB, table string) (map[string]bool, error) {
	rows, err := db.Query(`PRAGMA table_info(` + table + `)`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	columns := map[string]bool{}
	for rows.Next() {
		var cid sql.NullInt64
		var name sql.NullString
		var typ sql.NullString
		var notNull sql.NullInt64
		var defaultValue sql.NullString
		var pk sql.NullInt64
		if err := rows.Scan(&cid, &name, &typ, &notNull, &defaultValue, &pk); err != nil {
			return nil, err
		}
		if name.Valid {
			columns[name.String] = true
		}
	}
	return columns, rows.Err()
}

func textColumn(columns map[string]bool, name string) string {
	if !columns[name] {
		return "''"
	}
	return "COALESCE(CAST(" + name + " AS TEXT), '')"
}

func firstTextColumn(columns map[string]bool, names ...string) string {
	for _, name := range names {
		if columns[name] {
			return textColumn(columns, name)
		}
	}
	return "''"
}

func orderBy(columns map[string]bool, names ...string) string {
	var order []string
	for _, name := range names {
		if columns[name] {
			order = append(order, name)
		}
	}
	if len(order) == 0 {
		return "rowid"
	}
	return strings.Join(order, ", ")
}

func tableSource(kind string, sourcePath string, message string) SourceEvidence {
	return SourceEvidence{Kind: kind, Path: sourcePath, Confidence: "low", Message: message}
}

func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && !info.IsDir()
}

func isMissingTable(err error) bool {
	return err != nil && (strings.Contains(err.Error(), "no such table") || errors.Is(err, sql.ErrNoRows))
}