Skip to content

Checkpoints V2: Allow for a force flag to rerun a migration that previously completed #885

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
computermode merged 3 commits into from
Apr 10, 2026
+110 −24
Merged

Checkpoints V2: Allow for a force flag to rerun a migration that previously completed #885

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
efa5e1d
Allow user to force re-creating all v2 checkpoints during migration
computermode Apr 8, 2026
ba299c9
Merge branch 'transcript-start-at-metadata' of https://github.com/ent…
computermode Apr 10, 2026
a5d13d0
Merge branch 'main' into migrate-checkpoints-v2-force-option
computermode Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 15 additions & 9 deletions cmd/entire/cli/migrate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import (

func newMigrateCmd() *cobra.Command {
var checkpointsFlag string
var forceFlag bool

cmd := &cobra.Command{
Use: "migrate",
Expand Down Expand Up @@ -53,11 +54,12 @@ func newMigrateCmd() *cobra.Command {
} else {
defer logging.Close()
}
return runMigrateCheckpointsV2(ctx, cmd)
return runMigrateCheckpointsV2(ctx, cmd, forceFlag)
},
}

cmd.Flags().StringVar(&checkpointsFlag, "checkpoints", "", "Target checkpoint format version (e.g., \"v2\")")
cmd.Flags().BoolVar(&forceFlag, "force", false, "Force re-migration of all checkpoints, overwriting existing v2 data")

return cmd
}
Expand All @@ -68,7 +70,7 @@ type migrateResult struct {
failed int
}

func runMigrateCheckpointsV2(ctx context.Context, cmd *cobra.Command) error {
func runMigrateCheckpointsV2(ctx context.Context, cmd *cobra.Command, force bool) error {
repo, err := strategy.OpenRepository(ctx)
if err != nil {
cmd.SilenceUsage = true
Expand All @@ -80,7 +82,7 @@ func runMigrateCheckpointsV2(ctx context.Context, cmd *cobra.Command) error {
v2Store := checkpoint.NewV2GitStore(repo, migrateRemoteName)
out := cmd.OutOrStdout()

result, err := migrateCheckpointsV2(ctx, repo, v1Store, v2Store, out)
result, err := migrateCheckpointsV2(ctx, repo, v1Store, v2Store, out, force)
if err != nil {
return err
}
Expand All @@ -103,7 +105,7 @@ var (

const migrateRemoteName = "origin"

func migrateCheckpointsV2(ctx context.Context, repo *git.Repository, v1Store *checkpoint.GitStore, v2Store *checkpoint.V2GitStore, out io.Writer) (*migrateResult, error) {
func migrateCheckpointsV2(ctx context.Context, repo *git.Repository, v1Store *checkpoint.GitStore, v2Store *checkpoint.V2GitStore, out io.Writer, force bool) (*migrateResult, error) {
v1List, err := v1Store.ListCommitted(ctx)
if err != nil {
return nil, fmt.Errorf("failed to list v1 checkpoints: %w", err)
Expand All @@ -114,14 +116,18 @@ func migrateCheckpointsV2(ctx context.Context, repo *git.Repository, v1Store *ch
return &migrateResult{}, nil
}

fmt.Fprintln(out, "Migrating v1 checkpoints to v2...")
if force {
fmt.Fprintln(out, "Force-migrating v1 checkpoints to v2 (overwriting existing)...")
} else {
fmt.Fprintln(out, "Migrating v1 checkpoints to v2...")
}
total := len(v1List)
result := &migrateResult{}

for i, info := range v1List {
prefix := fmt.Sprintf(" [%d/%d] Migrating checkpoint %s...", i+1, total, info.CheckpointID)

if migrateErr := migrateOneCheckpoint(ctx, repo, v1Store, v2Store, info, out, prefix); migrateErr != nil {
if migrateErr := migrateOneCheckpoint(ctx, repo, v1Store, v2Store, info, out, prefix, force); migrateErr != nil {
switch {
case errors.Is(migrateErr, errAlreadyMigrated):
fmt.Fprintf(out, "%s skipped (already in v2)\n", prefix)
Expand All @@ -146,14 +152,14 @@ func migrateCheckpointsV2(ctx context.Context, repo *git.Repository, v1Store *ch
return result, nil
}

func migrateOneCheckpoint(ctx context.Context, repo *git.Repository, v1Store *checkpoint.GitStore, v2Store *checkpoint.V2GitStore, info checkpoint.CommittedInfo, out io.Writer, prefix string) error {
func migrateOneCheckpoint(ctx context.Context, repo *git.Repository, v1Store *checkpoint.GitStore, v2Store *checkpoint.V2GitStore, info checkpoint.CommittedInfo, out io.Writer, prefix string, force bool) error {
existing, err := v2Store.ReadCommitted(ctx, info.CheckpointID)
if err != nil {
return fmt.Errorf("failed to check v2 for checkpoint %s: %w", info.CheckpointID, err)
}

// Already in v2 — check if any aspect of sessions are missing and backfill
if existing != nil {
// Already in v2 — when not forcing, check if any aspect of sessions are missing and backfill
if existing != nil && !force {
repaired, repairErr := repairPartialV2Checkpoint(ctx, repo, v1Store, v2Store, info, existing)
if repairErr != nil {
return repairErr
Expand Down
110 changes: 95 additions & 15 deletions cmd/entire/cli/migrate_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func TestMigrateCheckpointsV2_Basic(t *testing.T) {

var stdout bytes.Buffer

result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 1, result.migrated)
assert.Equal(t, 0, result.skipped)
Expand All @@ -111,19 +111,99 @@ func TestMigrateCheckpointsV2_Idempotent(t *testing.T) {
var stdout bytes.Buffer

// First run: should migrate
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 1, result1.migrated)
assert.Equal(t, 0, result1.skipped)

// Second run: should skip (no agent type means backfill also can't produce compact transcript)
stdout.Reset()
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 0, result2.migrated)
assert.Equal(t, 1, result2.skipped)
}

func TestMigrateCheckpointsV2_ForceOverwritesExisting(t *testing.T) {
t.Parallel()
repo := initMigrateTestRepo(t)
v1Store, v2Store := newMigrateStores(repo)

cpID := id.MustCheckpointID("f0f1f2f3f4f5")
writeV1Checkpoint(t, v1Store, cpID, "session-force",
[]byte("{\"type\":\"assistant\",\"message\":\"original\"}\n"),
[]string{"original prompt"},
)

var stdout bytes.Buffer

// First run: normal migration
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 1, result1.migrated)

// Second run without force: should skip
stdout.Reset()
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 0, result2.migrated)
assert.Equal(t, 1, result2.skipped)

// Third run with force: should re-migrate
stdout.Reset()
result3, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, true)
require.NoError(t, err)
assert.Equal(t, 1, result3.migrated)
assert.Equal(t, 0, result3.skipped)
assert.Contains(t, stdout.String(), "Force-migrating")

// Verify checkpoint still readable in v2
summary, readErr := v2Store.ReadCommitted(context.Background(), cpID)
require.NoError(t, readErr)
require.NotNil(t, summary)
assert.Equal(t, cpID, summary.CheckpointID)
}

func TestMigrateCheckpointsV2_ForceMultipleCheckpoints(t *testing.T) {
t.Parallel()
repo := initMigrateTestRepo(t)
v1Store, v2Store := newMigrateStores(repo)

cpID1 := id.MustCheckpointID("a0a1a2a3a4a5")
cpID2 := id.MustCheckpointID("b0b1b2b3b4b5")
writeV1Checkpoint(t, v1Store, cpID1, "session-force-1",
[]byte("{\"type\":\"assistant\",\"message\":\"first\"}\n"),
[]string{"prompt 1"},
)
writeV1Checkpoint(t, v1Store, cpID2, "session-force-2",
[]byte("{\"type\":\"assistant\",\"message\":\"second\"}\n"),
[]string{"prompt 2"},
)

// First run: migrates both
var discard bytes.Buffer
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &discard, false)
require.NoError(t, err)
assert.Equal(t, 2, result1.migrated)

// Force re-migrate: should re-migrate both (0 skipped)
var stdout bytes.Buffer
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, true)
require.NoError(t, err)
assert.Equal(t, 2, result2.migrated)
assert.Equal(t, 0, result2.skipped)
}

func TestMigrateCmd_ForceFlag(t *testing.T) {
t.Parallel()
cmd := newMigrateCmd()

// Verify --force flag exists
flag := cmd.Flags().Lookup("force")
require.NotNil(t, flag, "--force flag should be registered")
assert.Equal(t, "false", flag.DefValue)
}

func TestMigrateCheckpointsV2_MultiSession(t *testing.T) {
t.Parallel()
repo := initMigrateTestRepo(t)
Expand All @@ -145,7 +225,7 @@ func TestMigrateCheckpointsV2_MultiSession(t *testing.T) {

var stdout bytes.Buffer

result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 1, result.migrated)

Expand All @@ -163,7 +243,7 @@ func TestMigrateCheckpointsV2_NoV1Branch(t *testing.T) {
var stdout bytes.Buffer

// No v1 data written — ListCommitted returns empty
result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 0, result.migrated)
assert.Contains(t, stdout.String(), "Nothing to migrate")
Expand Down Expand Up @@ -199,7 +279,7 @@ func TestMigrateCheckpointsV2_CompactionSkipped(t *testing.T) {

var stdout bytes.Buffer

result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, migrateErr)
assert.Equal(t, 1, result.migrated)
assert.Contains(t, stdout.String(), "compact transcript not generated")
Expand All @@ -226,7 +306,7 @@ func TestMigrateCheckpointsV2_TaskCheckpoint(t *testing.T) {

var stdout bytes.Buffer

result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, migrateErr)
assert.Equal(t, 1, result.migrated)

Expand Down Expand Up @@ -263,13 +343,13 @@ func TestMigrateCheckpointsV2_AllSkippedOnRerun(t *testing.T) {

// First run: migrates both
var discard bytes.Buffer
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &discard)
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &discard, false)
require.NoError(t, err)
assert.Equal(t, 2, result1.migrated)

// Second run: skips both
var stdout bytes.Buffer
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result2, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, err)
assert.Equal(t, 0, result2.migrated)
assert.Equal(t, 2, result2.skipped)
Expand Down Expand Up @@ -317,7 +397,7 @@ func TestMigrateCheckpointsV2_BackfillCompactTranscript(t *testing.T) {

// Run migration — should backfill the compact transcript
var stdout bytes.Buffer
result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout)
result, migrateErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &stdout, false)
require.NoError(t, migrateErr)
assert.Equal(t, 1, result.migrated, "backfill should count as migrated")
assert.Equal(t, 0, result.skipped)
Expand Down Expand Up @@ -372,7 +452,7 @@ func TestMigrateCheckpointsV2_UsesComputedCompactTranscriptStart(t *testing.T) {
require.Positive(t, expectedOffset, "expected non-zero compact transcript start")

var stdout bytes.Buffer
result, migrateErr := migrateCheckpointsV2(ctx, repo, v1Store, v2Store, &stdout)
result, migrateErr := migrateCheckpointsV2(ctx, repo, v1Store, v2Store, &stdout, false)
require.NoError(t, migrateErr)
assert.Equal(t, 1, result.migrated)

Expand Down Expand Up @@ -410,7 +490,7 @@ func TestMigrateCheckpointsV2_RepairsMissingFullTranscriptBeforeBackfill(t *test

// Initial migration to create v2 state.
var initialRun bytes.Buffer
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &initialRun)
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &initialRun, false)
require.NoError(t, err)
assert.Equal(t, 1, result1.migrated)

Expand All @@ -419,7 +499,7 @@ func TestMigrateCheckpointsV2_RepairsMissingFullTranscriptBeforeBackfill(t *test

// Re-run migration: should repair /full/current and count as migrated (not skipped).
var rerun bytes.Buffer
result2, rerunErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &rerun)
result2, rerunErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &rerun, false)
require.NoError(t, rerunErr)
assert.Equal(t, 1, result2.migrated)
assert.Equal(t, 0, result2.failed)
Expand All @@ -443,7 +523,7 @@ func TestMigrateCheckpointsV2_RepairsCurrentFullEvenWhenArchiveExists(t *testing

// Initial migration to seed v2.
var initialRun bytes.Buffer
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &initialRun)
result1, err := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &initialRun, false)
require.NoError(t, err)
assert.Equal(t, 1, result1.migrated)

Expand All @@ -463,7 +543,7 @@ func TestMigrateCheckpointsV2_RepairsCurrentFullEvenWhenArchiveExists(t *testing

// Re-run migration: should still repair /full/current.
var rerun bytes.Buffer
result2, rerunErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &rerun)
result2, rerunErr := migrateCheckpointsV2(context.Background(), repo, v1Store, v2Store, &rerun, false)
require.NoError(t, rerunErr)
assert.Equal(t, 1, result2.migrated)
assert.Contains(t, rerun.String(), "repaired partial v2 checkpoint state")
Expand Down
Loading