Merge pull request #1114 from entireio/more-migration-speedups

Checkpoints v2: Fix migration rerun logic (speedups, avoid duplicated migration efforts)

Author: computermode

cmd/entire/cli/checkpoint/v2_committed.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"strconv"
 	"strings"
 
 	"github.com/entireio/cli/cmd/entire/cli/agent"
@@ -135,6 +136,116 @@ func (s *V2GitStore) findFullSessionArtifacts(checkpointID id.CheckpointID, sess
 	return fullSessionArtifacts{}, nil
 }
 
+// FullSessionArtifactsIndex answers "does this session have complete /full/*
+// artifacts?" with an O(1) map lookup. Build it once via
+// BuildFullSessionArtifactsIndex.
+type FullSessionArtifactsIndex map[string]struct{}
+
+// Has reports whether the given session has a complete pair of
+// raw_transcript and raw_transcript_hash.txt entries in some /full/* ref.
+func (idx FullSessionArtifactsIndex) Has(checkpointID id.CheckpointID, sessionIndex int) bool {
+	if idx == nil {
+		return false
+	}
+	_, ok := idx[fullArtifactsIndexKey(checkpointID, sessionIndex)]
+	return ok
+}
+
+func fullArtifactsIndexKey(checkpointID id.CheckpointID, sessionIndex int) string {
+	return string(checkpointID) + "/" + strconv.Itoa(sessionIndex)
+}
+
+// BuildFullSessionArtifactsIndex walks every /full/* ref's tree once and
+// records sessions whose subtree contains both raw_transcript[/.NNN] and
+// raw_transcript_hash.txt. Amortizes per-session HasFullSessionArtifacts
+// calls — each of which would otherwise list every git ref and re-walk every
+// /full/* tree — across the rest of the run.
+func (s *V2GitStore) BuildFullSessionArtifactsIndex() (FullSessionArtifactsIndex, error) {
+	refNames, err := s.fullRefSearchOrder()
+	if err != nil {
+		return nil, err
+	}
+
+	index := make(FullSessionArtifactsIndex)
+	for _, refName := range refNames {
+		_, rootTreeHash, refErr := s.GetRefState(refName)
+		if refErr != nil {
+			if errors.Is(refErr, plumbing.ErrReferenceNotFound) {
+				continue
+			}
+			return nil, fmt.Errorf("read %s: %w", refName, refErr)
+		}
+		rootTree, treeErr := s.repo.TreeObject(rootTreeHash)
+		if treeErr != nil {
+			return nil, fmt.Errorf("read %s root tree: %w", refName, treeErr)
+		}
+		keys, err := s.listFullSessionsInTree(rootTree)
+		if err != nil {
+			return nil, fmt.Errorf("walk %s: %w", refName, err)
+		}
+		for _, key := range keys {
+			index[key] = struct{}{}
+		}
+	}
+	return index, nil
+}
+
+func (s *V2GitStore) listFullSessionsInTree(rootTree *object.Tree) ([]string, error) {
+	var keys []string
+	for _, shardEntry := range rootTree.Entries {
+		if shardEntry.Mode != filemode.Dir || len(shardEntry.Name) != 2 {
+			continue
+		}
+		shardTree, err := s.repo.TreeObject(shardEntry.Hash)
+		if err != nil {
+			return nil, fmt.Errorf("read shard %s: %w", shardEntry.Name, err)
+		}
+		for _, cpEntry := range shardTree.Entries {
+			if cpEntry.Mode != filemode.Dir {
+				continue
+			}
+			cpTree, err := s.repo.TreeObject(cpEntry.Hash)
+			if err != nil {
+				return nil, fmt.Errorf("read checkpoint tree %s/%s: %w", shardEntry.Name, cpEntry.Name, err)
+			}
+			cpid := id.CheckpointID(shardEntry.Name + cpEntry.Name)
+			for _, sessionEntry := range cpTree.Entries {
+				if sessionEntry.Mode != filemode.Dir {
+					continue
+				}
+				sessionIdx, atoiErr := strconv.Atoi(sessionEntry.Name)
+				if atoiErr != nil {
+					continue
+				}
+				sessionTree, err := s.repo.TreeObject(sessionEntry.Hash)
+				if err != nil {
+					return nil, fmt.Errorf("read session tree %s/%s/%d: %w", shardEntry.Name, cpEntry.Name, sessionIdx, err)
+				}
+				if !sessionHasCompleteFullArtifacts(sessionTree.Entries) {
+					continue
+				}
+				keys = append(keys, fullArtifactsIndexKey(cpid, sessionIdx))
+			}
+		}
+	}
+	return keys, nil
+}
+
+func sessionHasCompleteFullArtifacts(entries []object.TreeEntry) bool {
+	hasTranscript := false
+	hasHash := false
+	for _, entry := range entries {
+		switch {
+		case entry.Name == paths.V2RawTranscriptFileName,
+			strings.HasPrefix(entry.Name, paths.V2RawTranscriptFileName+"."):
+			hasTranscript = true
+		case entry.Name == paths.V2RawTranscriptHashFileName:
+			hasHash = true
+		}
+	}
+	return hasTranscript && hasHash
+}
+
 func (s *V2GitStore) fullRefSearchOrder() ([]plumbing.ReferenceName, error) {
 	refNames := []plumbing.ReferenceName{plumbing.ReferenceName(paths.V2FullCurrentRefName)}
 
@@ -861,57 +972,3 @@ func (s *V2GitStore) UpdateSummary(ctx context.Context, checkpointID id.Checkpoi
 	commitMsg := fmt.Sprintf("Update summary for checkpoint %s (session: %s)", checkpointID, metadata.SessionID)
 	return s.updateRef(ctx, refName, newTreeHash, parentHash, commitMsg, authorName, authorEmail)
 }
-
-// CleanupV1TranscriptFiles removes legacy v1-named transcript files (full.jsonl,
-// full.jsonl.*, content_hash.txt) from /full/current for a given checkpoint.
-// Older CLI versions wrote these before the rename to raw_transcript.
-// Returns nil if /full/current doesn't exist or no v1 files were found.
-func (s *V2GitStore) CleanupV1TranscriptFiles(ctx context.Context, checkpointID id.CheckpointID, sessionCount int) error {
-	refName := plumbing.ReferenceName(paths.V2FullCurrentRefName)
-	parentHash, rootTreeHash, err := s.GetRefState(refName)
-	if err != nil {
-		if errors.Is(err, plumbing.ErrReferenceNotFound) {
-			return nil // /full/current doesn't exist yet — nothing to clean
-		}
-		return err
-	}
-
-	checkpointPath := checkpointID.Path()
-	basePath := checkpointPath + "/"
-
-	entries, err := s.gs.flattenCheckpointEntries(rootTreeHash, checkpointPath)
-	if err != nil {
-		return err
-	}
-
-	changed := false
-	for sessionIdx := range sessionCount {
-		sessionPath := fmt.Sprintf("%s%d/", basePath, sessionIdx)
-		v1TranscriptPath := sessionPath + paths.TranscriptFileName
-		v1HashPath := sessionPath + paths.ContentHashFileName
-
-		for key := range entries {
-			switch {
-			case key == v1TranscriptPath,
-				strings.HasPrefix(key, v1TranscriptPath+"."),
-				key == v1HashPath:
-				delete(entries, key)
-				changed = true
-			}
-		}
-	}
-
-	if !changed {
-		return nil
-	}
-
-	newTreeHash, err := s.gs.spliceCheckpointSubtree(ctx, rootTreeHash, checkpointID, basePath, entries)
-	if err != nil {
-		return fmt.Errorf("tree surgery failed: %w", err)
-	}
-
-	authorName, authorEmail := GetGitAuthorFromRepo(s.repo)
-	return s.updateRef(ctx, refName, newTreeHash, parentHash,
-		fmt.Sprintf("Clean up v1 transcript files for %s\n", checkpointID),
-		authorName, authorEmail)
-}

cmd/entire/cli/checkpoint/v2_generation.go

@@ -333,6 +333,23 @@ func checkpointTimestampRangeFromFullTree(cpTree *object.Tree) (GenerationMetada
 	return gen, found
 }
 
+// AggregateTranscriptTimestamps derives a generation timestamp envelope from
+// transcripts already in memory, using the same first/last-event semantics
+// as ComputeGenerationTimestampsFromTrees but skipping the blob reads.
+func AggregateTranscriptTimestamps(transcripts [][]byte) (GenerationMetadata, bool) {
+	var gen GenerationMetadata
+	found := false
+	for _, transcript := range transcripts {
+		if len(transcript) == 0 {
+			continue
+		}
+		if r, ok := timestampRangeFromTranscript(transcript); ok {
+			mergeGenerationRange(&gen, &found, r)
+		}
+	}
+	return gen, found
+}
+
 func timestampRangeFromTranscript(transcript []byte) (GenerationMetadata, bool) {
 	reader := bufio.NewReader(bytes.NewReader(transcript))
 	var gen GenerationMetadata

cmd/entire/cli/checkpoint/v2_read.go

@@ -183,6 +183,53 @@ func (s *V2GitStore) ReadSessionCompactTranscript(ctx context.Context, checkpoin
 	return []byte(content), nil
 }
 
+// ReadSessionMetadata reads only the session's metadata.json from the v2
+// /main ref — cheaper than ReadSessionMetadataAndPrompts when the caller
+// only needs metadata fields (e.g. SessionID). Returns ErrCheckpointNotFound
+// if the checkpoint or session doesn't exist on /main.
+func (s *V2GitStore) ReadSessionMetadata(ctx context.Context, checkpointID id.CheckpointID, sessionIndex int) (CommittedMetadata, error) {
+	if err := ctx.Err(); err != nil {
+		return CommittedMetadata{}, err //nolint:wrapcheck // Propagating context cancellation
+	}
+
+	refName := plumbing.ReferenceName(paths.V2MainRefName)
+	_, rootTreeHash, err := s.GetRefState(refName)
+	if err != nil {
+		return CommittedMetadata{}, ErrCheckpointNotFound
+	}
+
+	rootTree, err := s.repo.TreeObject(rootTreeHash)
+	if err != nil {
+		return CommittedMetadata{}, ErrCheckpointNotFound
+	}
+
+	cpTree, err := rootTree.Tree(checkpointID.Path())
+	if err != nil {
+		return CommittedMetadata{}, ErrCheckpointNotFound
+	}
+
+	sessionTree, err := cpTree.Tree(strconv.Itoa(sessionIndex))
+	if err != nil {
+		return CommittedMetadata{}, ErrCheckpointNotFound
+	}
+
+	sessionFT := s.wrapWithFetcher(ctx, sessionTree)
+	metadataFile, err := sessionFT.File(paths.MetadataFileName)
+	if err != nil {
+		return CommittedMetadata{}, fmt.Errorf("read session metadata file: %w", err)
+	}
+	content, err := metadataFile.Contents()
+	if err != nil {
+		return CommittedMetadata{}, fmt.Errorf("read session metadata contents: %w", err)
+	}
+
+	var meta CommittedMetadata
+	if err := json.Unmarshal([]byte(content), &meta); err != nil {
+		return CommittedMetadata{}, fmt.Errorf("parse session metadata: %w", err)
+	}
+	return meta, nil
+}
+
 // ReadSessionMetadataAndPrompts reads a session's metadata and prompts from the
 // v2 /main ref without requiring the raw transcript from /full/* refs.
 // Used by explain when the raw transcript is unavailable but compact transcript

cmd/entire/cli/checkpoint/v2_read_test.go

@@ -2,6 +2,8 @@ package checkpoint
 
 import (
 	"context"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/entireio/cli/cmd/entire/cli/checkpoint/id"
@@ -10,6 +12,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/go-git/go-git/v6"
 	"github.com/go-git/go-git/v6/plumbing"
 	"github.com/go-git/go-git/v6/plumbing/filemode"
 	"github.com/go-git/go-git/v6/plumbing/object"
@@ -166,6 +169,75 @@ func TestV2ReadSessionMetadataAndPrompts_ReturnsWithoutTranscript(t *testing.T)
 	assert.Empty(t, content.Transcript)
 }
 
+func TestV2ReadSessionMetadata_ReturnsMetadata(t *testing.T) {
+	t.Parallel()
+	repo := initTestRepo(t)
+	store := NewV2GitStore(repo, "origin")
+	cpID := id.MustCheckpointID("f1f2f3f4f5fa")
+	ctx := context.Background()
+
+	err := store.WriteCommitted(ctx, WriteCommittedOptions{
+		CheckpointID: cpID,
+		SessionID:    "session-metadata-only",
+		Strategy:     "manual-commit",
+		Prompts:      []string{"test prompt"},
+		AuthorName:   "Test",
+		AuthorEmail:  "test@test.com",
+	})
+	require.NoError(t, err)
+
+	meta, err := store.ReadSessionMetadata(ctx, cpID, 0)
+	require.NoError(t, err)
+	assert.Equal(t, "session-metadata-only", meta.SessionID)
+}
+
+func TestV2ReadSessionMetadata_FetchesMissingMetadataBlob(t *testing.T) {
+	t.Parallel()
+	repo := initTestRepo(t)
+	store := NewV2GitStore(repo, "origin")
+	cpID := id.MustCheckpointID("f1f2f3f4f5fb")
+	ctx := context.Background()
+
+	err := store.WriteCommitted(ctx, WriteCommittedOptions{
+		CheckpointID: cpID,
+		SessionID:    "session-fetch-metadata",
+		Strategy:     "manual-commit",
+		Prompts:      []string{"test prompt"},
+		AuthorName:   "Test",
+		AuthorEmail:  "test@test.com",
+	})
+	require.NoError(t, err)
+
+	wt, err := repo.Worktree()
+	require.NoError(t, err)
+	repoRoot := wt.Filesystem.Root()
+	mainTree := v2MainTree(t, repo)
+	sessionTree, err := mainTree.Tree(cpID.Path() + "/0")
+	require.NoError(t, err)
+	metadataEntry, err := sessionTree.FindEntry(paths.MetadataFileName)
+	require.NoError(t, err)
+	metadataContent := v2ReadFile(t, mainTree, cpID.Path()+"/0/"+paths.MetadataFileName)
+
+	metadataObjectPath := filepath.Join(repoRoot, ".git", "objects", metadataEntry.Hash.String()[:2], metadataEntry.Hash.String()[2:])
+	require.NoError(t, os.Remove(metadataObjectPath))
+
+	reopenedRepo, err := git.PlainOpen(repoRoot)
+	require.NoError(t, err)
+	reopenedStore := NewV2GitStore(reopenedRepo, "origin")
+	fetchCalled := false
+	reopenedStore.SetBlobFetcher(func(_ context.Context, hashes []plumbing.Hash) error {
+		fetchCalled = true
+		require.Equal(t, []plumbing.Hash{metadataEntry.Hash}, hashes)
+		_, createErr := CreateBlobFromContent(reopenedRepo, []byte(metadataContent))
+		return createErr
+	})
+
+	meta, err := reopenedStore.ReadSessionMetadata(ctx, cpID, 0)
+	require.NoError(t, err)
+	assert.True(t, fetchCalled)
+	assert.Equal(t, "session-fetch-metadata", meta.SessionID)
+}
+
 func TestV2ReadSessionMetadataAndPrompts_MissingCheckpoint(t *testing.T) {
 	t.Parallel()
 	repo := initTestRepo(t)