Integrate chunk plan caching into the runner

internal/framework/checkpoint/loader.go

@@ -58,36 +58,6 @@ func (l *WorkspaceLoader) Source(moduleKey string) (pipeline.SourceCheckpoint, p
 	return pipeline.SourceCheckpoint{Document: &doc}, reusedDecision()
 }
 
-func (l *WorkspaceLoader) Chunk(moduleKey string, sourceDigest string) (pipeline.ChunkCheckpoint, pipeline.CheckpointDecision) {
-	expectedDependencies := digestFingerprints("source_document", sourceDigest)
-	var manifest coreworkspace.ChunkManifest
-	if decision := l.readJSON("chunk/manifest.json", &manifest); !decision.Reused {
-		return pipeline.ChunkCheckpoint{}, decision
-	}
-	if decision := l.validateManifest(manifest.StageManifest, coreworkspace.StageChunk, "", moduleKey, coreworkspace.StatusSucceeded, expectedDependencies); !decision.Reused {
-		return pipeline.ChunkCheckpoint{}, decision
-	}
-	var payload chunksEnvelope
-	if decision := l.readJSON("chunk/chunks.json", &payload); !decision.Reused {
-		return pipeline.ChunkCheckpoint{}, decision
-	}
-	chunks, err := sourceChunksFromEnvelope(payload.Chunks)
-	if err != nil {
-		return pipeline.ChunkCheckpoint{}, invalidDecision("chunk checkpoint payload is invalid: %v", err)
-	}
-	if len(chunks) == 0 {
-		return pipeline.ChunkCheckpoint{}, invalidDecision("chunk checkpoint payload has no chunks")
-	}
-	outputDigests, err := chunkOutputDigests(chunks)
-	if err != nil {
-		return pipeline.ChunkCheckpoint{}, invalidDecision("chunk checkpoint output cannot be digested: %v", err)
-	}
-	if !fingerprintsEqual(coreworkspaceToPipelineFingerprints(manifest.OutputDigests), outputDigests) {
-		return pipeline.ChunkCheckpoint{}, invalidDecision("chunk checkpoint output digests do not match payload")
-	}
-	return pipeline.ChunkCheckpoint{Chunks: chunks, Warnings: cloneWarnings(payload.Warnings)}, reusedDecision()
-}
-
 func (l *WorkspaceLoader) Extract(laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint) (pipeline.ExtractCheckpoint, pipeline.CheckpointDecision) {
 	var manifest coreworkspace.ExtractLaneManifest
 	if d := l.readJSON(laneManifestPath("extract", laneID), &manifest); !d.Reused {
@@ -232,40 +202,6 @@ func (l *WorkspaceLoader) validateLaneManifest(manifest coreworkspace.StageManif
 	return reusedDecision()
 }
 
-func sourceChunksFromEnvelope(values []chunkEnvelope) ([]source.Chunk, error) {
-	if len(values) == 0 {
-		return nil, nil
-	}
-	out := make([]source.Chunk, 0, len(values))
-	for _, value := range values {
-		content, err := contentFromEnvelope(value.Content)
-		if err != nil {
-			return nil, err
-		}
-		annotations, err := source.CanonicalizeChunkAnnotations(value.Annotations)
-		if err != nil {
-			return nil, fmt.Errorf("canonicalize chunk %q annotations: %w", value.ID, err)
-		}
-		planAnnotations, err := source.CanonicalizeChunkAnnotations(value.PlanAnnotations)
-		if err != nil {
-			return nil, fmt.Errorf("canonicalize chunk %q plan_annotations: %w", value.ID, err)
-		}
-		out = append(out, source.Chunk{
-			ID:              value.ID,
-			SourceID:        value.SourceID,
-			Index:           value.Index,
-			Ref:             value.Ref,
-			Content:         content,
-			MediaType:       value.Content.MediaType,
-			Units:           cloneSourceUnits(value.Units),
-			Metadata:        cloneMetadata(value.Metadata),
-			Annotations:     annotations,
-			PlanAnnotations: planAnnotations,
-		})
-	}
-	return out, nil
-}
-
 func contentFromEnvelope(value binaryEnvelope) ([]byte, error) {
 	content, err := base64.StdEncoding.DecodeString(value.ContentBase64)
 	if err != nil {

internal/framework/checkpoint/recorder.go

@@ -65,54 +65,6 @@ func (r *WorkspaceRecorder) SourceFailed(moduleKey string, err error) error {
 	return r.writeManifest("source/manifest.json", coreworkspace.SourceManifest{StageManifest: manifest})
 }
 
-func (r *WorkspaceRecorder) ChunkRunning(moduleKey string, sourceDigest string) error {
-	manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusRunning)
-	manifest.ModuleKey = moduleKey
-	manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
-	manifest.StartedAt = timePtr(r.timestamp())
-	return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
-}
-
-func (r *WorkspaceRecorder) ChunkSucceeded(moduleKey string, sourceDigest string, chunks []source.Chunk, warnings []contracts.Warning) error {
-	outputDigests, err := chunkOutputDigests(chunks)
-	if err != nil {
-		return fmt.Errorf("digest chunk checkpoint output: %w", err)
-	}
-	payload := chunksEnvelope{Chunks: chunkEnvelopes(chunks), Warnings: cloneWarnings(warnings)}
-	if err := r.writePayload("chunk/chunks.json", payload); err != nil {
-		return err
-	}
-	manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusSucceeded)
-	manifest.ModuleKey = moduleKey
-	manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
-	manifest.OutputDigests = workspaceFingerprints(outputDigests)
-	manifest.ValidationStatus = validationStatusString(warnings, nil)
-	manifest.CompletedAt = timePtr(r.timestamp())
-	return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{
-		StageManifest: manifest,
-		ChunkCount:    len(chunks),
-	})
-}
-
-func (r *WorkspaceRecorder) ChunkRejected(moduleKey string, sourceDigest string, rejected contracts.RejectedOutput) error {
-	manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusSucceededWithRejections)
-	manifest.ModuleKey = moduleKey
-	manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
-	manifest.ValidationStatus = "rejected"
-	manifest.Rejections = rejectionSummaries([]contracts.RejectedOutput{rejected})
-	manifest.CompletedAt = timePtr(r.timestamp())
-	return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
-}
-
-func (r *WorkspaceRecorder) ChunkFailed(moduleKey string, sourceDigest string, err error) error {
-	manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusFailed)
-	manifest.ModuleKey = moduleKey
-	manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
-	manifest.CompletedAt = timePtr(r.timestamp())
-	manifest.Metadata = errorMetadata(err)
-	return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
-}
-
 func (r *WorkspaceRecorder) ExtractRunning(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
 	manifest := r.laneManifest(coreworkspace.StageExtract, coreworkspace.StatusRunning, laneID, moduleKey, dependencies)
 	manifest.StartedAt = timePtr(r.timestamp())
@@ -245,23 +197,6 @@ type sourceDocumentEnvelope struct {
 	Document source.SourceDocument `json:"document"`
 }
 
-type chunksEnvelope struct {
-	Chunks   []chunkEnvelope     `json:"chunks"`
-	Warnings []contracts.Warning `json:"warnings,omitempty"`
-}
-
-type chunkEnvelope struct {
-	ID              string                  `json:"id"`
-	SourceID        string                  `json:"source_id"`
-	Index           int                     `json:"index"`
-	Ref             source.SourceRef        `json:"ref"`
-	Content         binaryEnvelope          `json:"content"`
-	Units           []source.SourceUnit     `json:"units,omitempty"`
-	Metadata        map[string]any          `json:"metadata,omitempty"`
-	Annotations     source.ChunkAnnotations `json:"annotations,omitempty"`
-	PlanAnnotations source.ChunkAnnotations `json:"plan_annotations,omitempty"`
-}
-
 type binaryEnvelope struct {
 	ContentBase64 string              `json:"content_base64,omitempty"`
 	ContentDigest string              `json:"content_digest,omitempty"`
@@ -315,27 +250,6 @@ func artifactOutputDigests(outputs []pipeline.CheckpointArtifact) []pipeline.Che
 	return normalizeFingerprints(values)
 }
 
-func chunkEnvelopes(chunks []source.Chunk) []chunkEnvelope {
-	if len(chunks) == 0 {
-		return nil
-	}
-	out := make([]chunkEnvelope, 0, len(chunks))
-	for _, chunk := range chunks {
-		out = append(out, chunkEnvelope{
-			ID:              chunk.ID,
-			SourceID:        chunk.SourceID,
-			Index:           chunk.Index,
-			Ref:             chunk.Ref,
-			Content:         binaryEnvelopeFromContent(chunk.Content, chunk.MediaType, chunk.Metadata, nil),
-			Units:           cloneSourceUnits(chunk.Units),
-			Metadata:        cloneMetadata(chunk.Metadata),
-			Annotations:     source.CloneChunkAnnotations(chunk.Annotations),
-			PlanAnnotations: source.CloneChunkAnnotations(chunk.PlanAnnotations),
-		})
-	}
-	return out
-}
-
 func binaryEnvelopeFromContent(content []byte, mediaType string, metadata map[string]any, warnings []contracts.Warning) binaryEnvelope {
 	return binaryEnvelope{
 		ContentBase64: base64.StdEncoding.EncodeToString(content),
@@ -394,21 +308,6 @@ func cloneMetadata(metadata map[string]any) map[string]any {
 	return out
 }
 
-func chunkOutputDigests(chunks []source.Chunk) ([]pipeline.CheckpointFingerprint, error) {
-	values := make([]pipeline.CheckpointFingerprint, 0, len(chunks))
-	for _, chunk := range chunks {
-		digest, err := source.DigestChunk(chunk)
-		if err != nil {
-			return nil, fmt.Errorf("chunk %q: %w", chunk.ID, err)
-		}
-		values = append(values, pipeline.CheckpointFingerprint{
-			Name:  chunk.ID,
-			Value: digest,
-		})
-	}
-	return normalizeFingerprints(values), nil
-}
-
 func digestFingerprints(name string, digest string) []pipeline.CheckpointFingerprint {
 	digest = strings.TrimSpace(digest)
 	if digest == "" {

internal/framework/checkpoint/recorder_test.go

@@ -1,11 +1,9 @@
 package checkpoint
 
 import (
-	"encoding/base64"
 	"encoding/json"
 	"os"
 	"path/filepath"
-	"reflect"
 	"strings"
 	"testing"
 
@@ -25,24 +23,6 @@ func TestWorkspaceRecorderWritesSuccessfulCheckpointFiles(t *testing.T) {
 		Digest: "sha256:source",
 		Units:  []source.SourceUnit{{ID: 1, Kind: "line", Text: "hello", Ref: source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 1}}},
 	}
-	chunks := []source.Chunk{
-		{
-			ID:        "chunk-1",
-			SourceID:  "source-1",
-			Index:     0,
-			Ref:       source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 1},
-			Content:   []byte("chunk content"),
-			MediaType: "text/plain",
-			Units:     doc.Units,
-			Annotations: source.ChunkAnnotations{
-				"same": json.RawMessage(`{"range":1}`),
-			},
-			PlanAnnotations: source.ChunkAnnotations{
-				"same": json.RawMessage(`{"plan":2}`),
-			},
-		},
-	}
-
 	if err := recorder.SourceRunning("seriatim"); err != nil {
 		t.Fatalf("SourceRunning: %v", err)
 	}
@@ -54,60 +34,6 @@ func TestWorkspaceRecorderWritesSuccessfulCheckpointFiles(t *testing.T) {
 	if _, err := os.Stat(filepath.Join(root, "source", "source-document.json")); err != nil {
 		t.Fatalf("expected source checkpoint payload: %v", err)
 	}
-
-	if err := recorder.ChunkRunning("generic", doc.Digest); err != nil {
-		t.Fatalf("ChunkRunning: %v", err)
-	}
-	if err := recorder.ChunkSucceeded("generic", doc.Digest, chunks, nil); err != nil {
-		t.Fatalf("ChunkSucceeded: %v", err)
-	}
-	assertManifestStatus(t, filepath.Join(root, "chunk", "manifest.json"), coreworkspace.StatusSucceeded)
-	var chunkPayload struct {
-		Chunks []struct {
-			Annotations     source.ChunkAnnotations `json:"annotations"`
-			PlanAnnotations source.ChunkAnnotations `json:"plan_annotations"`
-			Content         struct {
-				ContentBase64 string `json:"content_base64"`
-				ContentDigest string `json:"content_digest"`
-			} `json:"content"`
-		} `json:"chunks"`
-	}
-	readJSON(t, filepath.Join(root, "chunk", "chunks.json"), &chunkPayload)
-	if len(chunkPayload.Chunks) != 1 {
-		t.Fatalf("checkpoint chunks = %#v, want one", chunkPayload.Chunks)
-	}
-	decoded, err := base64.StdEncoding.DecodeString(chunkPayload.Chunks[0].Content.ContentBase64)
-	if err != nil {
-		t.Fatalf("decode chunk content: %v", err)
-	}
-	if string(decoded) != "chunk content" {
-		t.Fatalf("chunk content = %q, want original content", decoded)
-	}
-	if got, want := chunkPayload.Chunks[0].Content.ContentDigest, contentDigest([]byte("chunk content")); got != want {
-		t.Fatalf("content digest = %q, want %q", got, want)
-	}
-	if !jsonEqual(t, chunkPayload.Chunks[0].Annotations["same"], json.RawMessage(`{"range":1}`)) || !jsonEqual(t, chunkPayload.Chunks[0].PlanAnnotations["same"], json.RawMessage(`{"plan":2}`)) {
-		t.Fatalf("checkpoint annotations = %s / %s", chunkPayload.Chunks[0].Annotations["same"], chunkPayload.Chunks[0].PlanAnnotations["same"])
-	}
-	loaded, decision := (&WorkspaceLoader{root: root}).Chunk("generic", doc.Digest)
-	if !decision.Reused || len(loaded.Chunks) != 1 {
-		t.Fatalf("loaded chunk checkpoint = %#v, decision = %#v", loaded, decision)
-	}
-	if string(loaded.Chunks[0].Annotations["same"]) != `{"range":1}` || string(loaded.Chunks[0].PlanAnnotations["same"]) != `{"plan":2}` {
-		t.Fatalf("loaded canonical annotations = %s / %s", loaded.Chunks[0].Annotations["same"], loaded.Chunks[0].PlanAnnotations["same"])
-	}
-}
-
-func jsonEqual(t *testing.T, left, right []byte) bool {
-	t.Helper()
-	var leftValue, rightValue any
-	if err := json.Unmarshal(left, &leftValue); err != nil {
-		t.Fatalf("decode left JSON: %v", err)
-	}
-	if err := json.Unmarshal(right, &rightValue); err != nil {
-		t.Fatalf("decode right JSON: %v", err)
-	}
-	return reflect.DeepEqual(leftValue, rightValue)
 }
 
 func TestWorkspaceArtifactCheckpointsRoundTripCodecIdentityAndBytes(t *testing.T) {
@@ -158,24 +84,6 @@ func TestWorkspaceLoaderInvalidatesMissingCorruptAndMismatchedCheckpoints(t *tes
 		}
 	})
 
-	t.Run("dependency mismatch", func(t *testing.T) {
-		root := t.TempDir()
-		recorder := newTestRecorder(t, root)
-		chunks := []source.Chunk{{
-			ID:        "chunk-1",
-			SourceID:  "source-1",
-			Content:   []byte("chunk content"),
-			MediaType: "text/plain",
-		}}
-		if err := recorder.ChunkSucceeded("generic", "sha256:source-a", chunks, nil); err != nil {
-			t.Fatalf("ChunkSucceeded: %v", err)
-		}
-		loader := &WorkspaceLoader{root: root}
-		if _, decision := loader.Chunk("generic", "sha256:source-b"); decision.Reused || !strings.Contains(decision.Reason, "dependency") {
-			t.Fatalf("decision = %#v, want dependency invalidation", decision)
-		}
-	})
-
 	t.Run("incompatible workspace schema remains untouched", func(t *testing.T) {
 		root := t.TempDir()
 		recorder := newTestRecorder(t, root)
@@ -207,28 +115,41 @@ func TestWorkspaceLoaderInvalidatesMissingCorruptAndMismatchedCheckpoints(t *tes
 		}
 	})
 
-	t.Run("corrupt payload", func(t *testing.T) {
-		root := t.TempDir()
-		recorder := newTestRecorder(t, root)
-		chunks := []source.Chunk{{
-			ID:        "chunk-1",
-			SourceID:  "source-1",
-			Content:   []byte("chunk content"),
-			MediaType: "text/plain",
-		}}
-		if err := recorder.ChunkSucceeded("generic", "sha256:source", chunks, nil); err != nil {
-			t.Fatalf("ChunkSucceeded: %v", err)
-		}
-		payloadPath := filepath.Join(root, "chunk", "chunks.json")
-		data := strings.ReplaceAll(string(readFile(t, payloadPath)), contentDigest([]byte("chunk content")), "sha256:bad")
-		if err := os.WriteFile(payloadPath, []byte(data), 0o644); err != nil {
-			t.Fatalf("corrupt chunk payload: %v", err)
-		}
-		loader := &WorkspaceLoader{root: root}
-		if _, decision := loader.Chunk("generic", "sha256:source"); decision.Reused || !strings.Contains(decision.Reason, "invalid") {
-			t.Fatalf("decision = %#v, want corrupt payload invalidation", decision)
-		}
-	})
+}
+
+func TestWorkspaceCheckpointsIgnoreLegacyChunkFiles(t *testing.T) {
+	root := t.TempDir()
+	legacyManifest := []byte(`{"legacy":"manifest"}`)
+	legacyPayload := []byte(`{"legacy":"chunks"}`)
+	legacyDir := filepath.Join(root, "chunk")
+	if err := os.MkdirAll(legacyDir, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	manifestPath := filepath.Join(legacyDir, "manifest.json")
+	payloadPath := filepath.Join(legacyDir, "chunks.json")
+	if err := os.WriteFile(manifestPath, legacyManifest, 0o600); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(payloadPath, legacyPayload, 0o600); err != nil {
+		t.Fatal(err)
+	}
+
+	doc := &source.SourceDocument{ID: "source-1", Kind: "document", Format: "text/plain", Units: []source.SourceUnit{{ID: 1, Kind: "line", Text: "hello", Ref: source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 1}}}}
+	doc.Digest, _ = source.DigestDocument(doc)
+	recorder := newTestRecorder(t, root)
+	if err := recorder.SourceSucceeded("seriatim", doc); err != nil {
+		t.Fatal(err)
+	}
+	loaded, decision := (&WorkspaceLoader{root: root}).Source("seriatim")
+	if !decision.Reused || loaded.Document == nil || loaded.Document.Digest != doc.Digest {
+		t.Fatalf("source checkpoint = %#v decision = %#v", loaded, decision)
+	}
+	if got := readFile(t, manifestPath); string(got) != string(legacyManifest) {
+		t.Fatalf("legacy manifest changed: %s", got)
+	}
+	if got := readFile(t, payloadPath); string(got) != string(legacyPayload) {
+		t.Fatalf("legacy payload changed: %s", got)
+	}
 }
 
 func TestWorkspaceRecorderRecordsRejectedExtractOutputs(t *testing.T) {