Add pipeline stage contracts

internal/framework/contracts/composition_test.go

@@ -12,14 +12,22 @@ import (
 )
 
 var _ contracts.InputAdapter = compositionAdapter{}
+var _ contracts.Chunker = compositionChunker{}
 var _ contracts.Extractor = compositionExtractor{}
+var _ contracts.Merger = compositionMerger{}
+var _ contracts.Normalizer = compositionNormalizer{}
 var _ contracts.Validator = compositionValidator{}
+var _ contracts.OutputEncoder = compositionOutputEncoder{}
 
 func TestContractsComposeAcrossPackages(t *testing.T) {
 	ctx := context.Background()
 	adapter := compositionAdapter{}
+	chunker := compositionChunker{}
 	extractor := compositionExtractor{}
+	merger := compositionMerger{}
+	normalizer := compositionNormalizer{}
 	validator := compositionValidator{}
+	encoder := compositionOutputEncoder{}
 
 	doc, err := adapter.Parse(ctx, contracts.ParseRequest{SourceID: "source-1"})
 	if err != nil {
@@ -29,7 +37,22 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
 		t.Fatalf("ValidateDocument() error = %v, want nil", err)
 	}
 
-	extraction, err := extractor.Extract(ctx, contracts.ExtractionRequest{Source: doc})
+	chunking, err := chunker.Chunk(ctx, contracts.ChunkRequest{
+		Source:   doc,
+		Metadata: map[string]any{"max_units": 2},
+	})
+	if err != nil {
+		t.Fatalf("Chunk() error = %v, want nil", err)
+	}
+	if len(chunking.Chunks) != 1 {
+		t.Fatalf("len(Chunks) = %d, want 1", len(chunking.Chunks))
+	}
+
+	extraction, err := extractor.Extract(ctx, contracts.ExtractionRequest{
+		Source:         doc,
+		Chunk:          &chunking.Chunks[0],
+		AmbientContext: map[string]any{"synopsis": "example synopsis"},
+	})
 	if err != nil {
 		t.Fatalf("Extract() error = %v, want nil", err)
 	}
@@ -44,9 +67,38 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
 		}
 	}
 
+	merge, err := merger.Merge(ctx, contracts.MergeRequest{
+		Source: doc,
+		LaneID: candidate.ArtifactType,
+		ChunkArtifacts: []contracts.ChunkArtifacts{
+			{
+				Chunk:      chunking.Chunks[0],
+				Candidates: extraction.Candidates,
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Merge() error = %v, want nil", err)
+	}
+	if len(merge.Candidates) != 1 {
+		t.Fatalf("len(merge.Candidates) = %d, want 1", len(merge.Candidates))
+	}
+
+	normalize, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
+		Source:     doc,
+		LaneID:     candidate.ArtifactType,
+		Candidates: merge.Candidates,
+	})
+	if err != nil {
+		t.Fatalf("Normalize() error = %v, want nil", err)
+	}
+	if len(normalize.Candidates) != 1 {
+		t.Fatalf("len(normalize.Candidates) = %d, want 1", len(normalize.Candidates))
+	}
+
 	validation, err := validator.Validate(ctx, contracts.ValidationRequest{
 		Source:     doc,
-		Candidates: extraction.Candidates,
+		Candidates: normalize.Candidates,
 	})
 	if err != nil {
 		t.Fatalf("Validate() error = %v, want nil", err)
@@ -62,6 +114,22 @@ func TestContractsComposeAcrossPackages(t *testing.T) {
 	if decision.CandidateIndex != candidate.Index {
 		t.Fatalf("CandidateIndex = %d, want %d", decision.CandidateIndex, candidate.Index)
 	}
+
+	output, err := encoder.Encode(ctx, contracts.OutputRequest{
+		Manifest: artifacts.RunManifest{RunID: "run-1"},
+		Approved: []artifacts.Artifact{
+			artifacts.ArtifactFromCandidate(normalize.Candidates[0]),
+		},
+	})
+	if err != nil {
+		t.Fatalf("Encode() error = %v, want nil", err)
+	}
+	if output.ContentType != "application/json" {
+		t.Fatalf("ContentType = %q, want application/json", output.ContentType)
+	}
+	if len(output.Bytes) == 0 {
+		t.Fatal("len(Bytes) = 0, want encoded bytes")
+	}
 }
 
 type compositionAdapter struct{}
@@ -83,6 +151,30 @@ func (adapter compositionAdapter) Parse(ctx context.Context, req contracts.Parse
 	}, nil
 }
 
+type compositionChunker struct{}
+
+func (chunker compositionChunker) Key() string {
+	return "generic-chunker"
+}
+
+func (chunker compositionChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
+	if req.Source == nil {
+		return contracts.ChunkResult{}, errors.New("source document is required")
+	}
+
+	return contracts.ChunkResult{
+		Chunks: []contracts.SourceChunk{
+			{
+				ID:       req.Source.ID + ":chunk:0",
+				SourceID: req.Source.ID,
+				Index:    0,
+				Units:    append([]source.SourceUnit(nil), req.Source.Units...),
+				Metadata: map[string]any{"strategy": "whole-document"},
+			},
+		},
+	}, nil
+}
+
 type compositionExtractor struct{}
 
 func (extractor compositionExtractor) Key() string {
@@ -105,6 +197,13 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
 	if req.Source == nil {
 		return contracts.ExtractionResult{}, errors.New("source document is required")
 	}
+	units := req.Source.Units
+	if req.Chunk != nil {
+		units = req.Chunk.Units
+	}
+	if req.AmbientContext["synopsis"] == "" {
+		return contracts.ExtractionResult{}, errors.New("ambient synopsis is required")
+	}
 
 	return contracts.ExtractionResult{
 		Candidates: []artifacts.ArtifactCandidate{
@@ -117,8 +216,8 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
 				SourceRefs: []source.SourceRef{
 					{
 						SourceID:    req.Source.ID,
-						StartUnitID: req.Source.Units[0].ID,
-						EndUnitID:   req.Source.Units[1].ID,
+						StartUnitID: units[0].ID,
+						EndUnitID:   units[len(units)-1].ID,
 					},
 				},
 			},
@@ -126,6 +225,31 @@ func (extractor compositionExtractor) Extract(ctx context.Context, req contracts
 	}, nil
 }
 
+type compositionMerger struct{}
+
+func (merger compositionMerger) Key() string {
+	return "generic-merger"
+}
+
+func (merger compositionMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
+	var candidates []artifacts.ArtifactCandidate
+	for _, chunkArtifacts := range req.ChunkArtifacts {
+		candidates = append(candidates, chunkArtifacts.Candidates...)
+	}
+
+	return contracts.MergeResult{Candidates: candidates}, nil
+}
+
+type compositionNormalizer struct{}
+
+func (normalizer compositionNormalizer) Key() string {
+	return "generic-normalizer"
+}
+
+func (normalizer compositionNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
+	return contracts.NormalizeResult{Candidates: req.Candidates}, nil
+}
+
 type compositionValidator struct{}
 
 func (validator compositionValidator) Name() string {
@@ -148,3 +272,28 @@ func (validator compositionValidator) Validate(ctx context.Context, req contract
 		Decisions:     decisions,
 	}, nil
 }
+
+type compositionOutputEncoder struct{}
+
+func (encoder compositionOutputEncoder) Key() string {
+	return "generic-output"
+}
+
+func (encoder compositionOutputEncoder) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
+	payload := struct {
+		RunID         string `json:"run_id"`
+		ApprovedCount int    `json:"approved_count"`
+	}{
+		RunID:         req.Manifest.RunID,
+		ApprovedCount: len(req.Approved),
+	}
+	encoded, err := json.Marshal(payload)
+	if err != nil {
+		return contracts.OutputResult{}, err
+	}
+
+	return contracts.OutputResult{
+		Bytes:       encoded,
+		ContentType: "application/json",
+	}, nil
+}