diff --git a/cmd/entire/cli/migrate.go b/cmd/entire/cli/migrate.go index 5a610560d..12fb9ab5c 100644 --- a/cmd/entire/cli/migrate.go +++ b/cmd/entire/cli/migrate.go @@ -25,6 +25,7 @@ import ( func newMigrateCmd() *cobra.Command { var checkpointsFlag string + var forceFlag bool cmd := &cobra.Command{ Use: "migrate", @@ -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 } @@ -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 @@ -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 } @@ -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) @@ -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) @@ -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 diff --git a/cmd/entire/cli/migrate_test.go b/cmd/entire/cli/migrate_test.go index edd16d31a..29ef2394d 100644 --- a/cmd/entire/cli/migrate_test.go +++ b/cmd/entire/cli/migrate_test.go @@ -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) @@ -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) @@ -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) @@ -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") @@ -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") @@ -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) @@ -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) @@ -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) @@ -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) @@ -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) @@ -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) @@ -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) @@ -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")