iop/packages/go/singlerequesttemplate/template.go

package singlerequesttemplate

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"regexp"
	"strings"
)

const MaxTemplateBytes = 8192

const DefaultPlanTemplate = `# Plan

## Goal
{{goal}}

## Steps
{{steps}}

## Verification
{{verification}}
`

const DefaultReviewTemplate = `# Review

## Worker Item Status
{{item_status}}

## Worker Changes
{{changes}}

## Worker Verification
{{verification}}

## Deviations
{{deviations}}
`

var (
	ErrInvalidTemplate  = errors.New("single-request template: invalid template")
	ErrTemplateTooLarge = errors.New("single-request template: exceeds maximum size")
	ErrMalformedPlan    = errors.New("single-request template: malformed plan output")
	ErrMalformedReview  = errors.New("single-request template: malformed review fields")
)

var placeholderRegex = regexp.MustCompile(`\{\{[^}]*\}\}`)

var (
	planPlaceholders   = []string{"{{goal}}", "{{steps}}", "{{verification}}"}
	planHeadings       = []string{"# Plan", "## Goal", "## Steps", "## Verification"}
	reviewPlaceholders = []string{"{{item_status}}", "{{changes}}", "{{verification}}", "{{deviations}}"}
	reviewHeadings     = []string{"# Review", "## Worker Item Status", "## Worker Changes", "## Worker Verification", "## Deviations"}
)

type ReviewFields struct {
	ItemStatus   string
	Changes      string
	Verification string
	Deviations   string
}

type PlanFields struct {
	Goal         string
	Steps        []string
	Verification []string
}

// exactLineOffsets returns the byte offset of the first standalone line equal to
// want and how many standalone lines matched. A standalone line is a maximal
// "\n"-delimited segment compared byte for byte, so decorated variants such as
// "### Plan", "# Plan Mismatch", or "NOTPASS" never satisfy a required line.
func exactLineOffsets(tmpl, want string) (int, int) {
	index, count, offset := -1, 0, 0
	for {
		var line string
		end := strings.IndexByte(tmpl[offset:], '\n')
		if end < 0 {
			line = tmpl[offset:]
		} else {
			line = tmpl[offset : offset+end]
		}
		if line == want {
			count++
			if index < 0 {
				index = offset
			}
		}
		if end < 0 {
			return index, count
		}
		offset += end + 1
	}
}

// requireExactLines resolves every required standalone line, rejecting missing
// and duplicated occurrences, and returns their offsets in the requested order.
// Only the documented required line is echoed on failure; template content is
// never included in the error.
func requireExactLines(tmpl string, lines []string) ([]int, error) {
	offsets := make([]int, len(lines))
	for i, line := range lines {
		index, count := exactLineOffsets(tmpl, line)
		if count != 1 {
			return nil, fmt.Errorf("%w: must contain the standalone line %q exactly once", ErrInvalidTemplate, line)
		}
		offsets[i] = index
	}
	return offsets, nil
}

// requirePlaceholderInventory closes the placeholder grammar. Each documented
// placeholder must appear exactly once; after removing those exact occurrences
// no template delimiter may survive, so unknown placeholders and unbalanced
// "{{"/"}}" residue are both rejected. Only documented placeholder names are
// echoed on failure; template content is never included in the error.
func requirePlaceholderInventory(tmpl string, placeholders []string) ([]int, error) {
	offsets := make([]int, len(placeholders))
	residue := tmpl
	for i, placeholder := range placeholders {
		if strings.Count(tmpl, placeholder) != 1 {
			return nil, fmt.Errorf("%w: must contain %s exactly once", ErrInvalidTemplate, placeholder)
		}
		offsets[i] = strings.Index(tmpl, placeholder)
		residue = strings.Replace(residue, placeholder, "", 1)
	}
	if placeholderRegex.MatchString(residue) {
		return nil, fmt.Errorf("%w: template declares an unknown placeholder", ErrInvalidTemplate)
	}
	if strings.Contains(residue, "{{") || strings.Contains(residue, "}}") {
		return nil, fmt.Errorf("%w: template leaves an unbalanced placeholder delimiter", ErrInvalidTemplate)
	}
	return offsets, nil
}

func ascending(values ...int) bool {
	for i := 1; i < len(values); i++ {
		if values[i-1] >= values[i] {
			return false
		}
	}
	return true
}

func Digest(content string) string {
	h := sha256.Sum256([]byte(content))
	return hex.EncodeToString(h[:])
}

func ValidatePlanTemplate(tmpl string) error {
	if len(tmpl) == 0 {
		return fmt.Errorf("%w: template is empty", ErrInvalidTemplate)
	}
	if len(tmpl) > MaxTemplateBytes {
		return fmt.Errorf("%w: template size %d exceeds max %d", ErrTemplateTooLarge, len(tmpl), MaxTemplateBytes)
	}

	placeholders, err := requirePlaceholderInventory(tmpl, planPlaceholders)
	if err != nil {
		return err
	}
	idxGoal, idxSteps, idxVerif := placeholders[0], placeholders[1], placeholders[2]
	if !ascending(idxGoal, idxSteps, idxVerif) {
		return fmt.Errorf("%w: placeholders must appear in order {{goal}}, {{steps}}, {{verification}}", ErrInvalidTemplate)
	}

	headings, err := requireExactLines(tmpl, planHeadings)
	if err != nil {
		return err
	}
	idxPlanH, idxGoalH, idxStepsH, idxVerifH := headings[0], headings[1], headings[2], headings[3]
	if !ascending(idxPlanH, idxGoalH, idxGoal, idxStepsH, idxSteps, idxVerifH, idxVerif) {
		return fmt.Errorf("%w: headings and placeholders must follow exact structural order", ErrInvalidTemplate)
	}

	return nil
}

func ValidateReviewTemplate(tmpl string) error {
	if len(tmpl) == 0 {
		return fmt.Errorf("%w: template is empty", ErrInvalidTemplate)
	}
	if len(tmpl) > MaxTemplateBytes {
		return fmt.Errorf("%w: template size %d exceeds max %d", ErrTemplateTooLarge, len(tmpl), MaxTemplateBytes)
	}

	placeholders, err := requirePlaceholderInventory(tmpl, reviewPlaceholders)
	if err != nil {
		return err
	}
	idxItemStatus, idxChanges, idxVerif, idxDeviations := placeholders[0], placeholders[1], placeholders[2], placeholders[3]
	if !ascending(idxItemStatus, idxChanges, idxVerif, idxDeviations) {
		return fmt.Errorf("%w: placeholders must appear in order {{item_status}}, {{changes}}, {{verification}}, {{deviations}}", ErrInvalidTemplate)
	}

	headings, err := requireExactLines(tmpl, reviewHeadings)
	if err != nil {
		return err
	}
	idxReviewH, idxItemStatusH, idxChangesH, idxVerifH, idxDeviationsH := headings[0], headings[1], headings[2], headings[3], headings[4]
	if !ascending(idxReviewH, idxItemStatusH, idxItemStatus, idxChangesH, idxChanges, idxVerifH, idxVerif, idxDeviationsH, idxDeviations) {
		return fmt.Errorf("%w: headings and placeholders must follow exact structural order", ErrInvalidTemplate)
	}

	// Close the heading set: only the documented worker headings may appear. A
	// reviewer-only section (Result, Checks, Summary) or any other markdown
	// heading would let the template describe a reviewer verdict or final review
	// page, so it is rejected rather than silently tolerated.
	if err := rejectUnknownMarkdownHeadings(tmpl, reviewHeadings); err != nil {
		return err
	}

	return nil
}

// rejectUnknownMarkdownHeadings ensures every standalone line beginning with
// "#" is one of the allowed documented headings. Decorated variants (e.g.
// "### Review") and reviewer-only headings ("## Result") are rejected because
// the only standalone matching already happens in requireExactLines; here we
// additionally forbid any extra heading that is not in the closed set.
func rejectUnknownMarkdownHeadings(tmpl string, allowed []string) error {
	allowedSet := make(map[string]struct{}, len(allowed))
	for _, line := range allowed {
		allowedSet[line] = struct{}{}
	}
	offset := 0
	for {
		var line string
		end := strings.IndexByte(tmpl[offset:], '\n')
		if end < 0 {
			line = tmpl[offset:]
		} else {
			line = tmpl[offset : offset+end]
		}
		if strings.HasPrefix(line, "#") {
			if _, ok := allowedSet[line]; !ok {
				return fmt.Errorf("%w: template declares an unknown heading", ErrInvalidTemplate)
			}
		}
		if end < 0 {
			return nil
		}
		offset += end + 1
	}
}

func normalizePlanFields(fields PlanFields) (string, string, string, error) {
	goal := strings.TrimSpace(fields.Goal)
	if goal == "" || strings.ContainsAny(goal, "\r\n") || strings.Contains(goal, "{{") || strings.Contains(goal, "}}") {
		return "", "", "", ErrMalformedPlan
	}
	normalizeStepItems := func(items []string, minimum, maximum int) (string, error) {
		if len(items) < minimum || len(items) > maximum {
			return "", ErrMalformedPlan
		}
		lines := make([]string, len(items))
		for i, item := range items {
			item = strings.TrimSpace(item)
			if item == "" || strings.ContainsAny(item, "\r\n") || strings.Contains(item, "{{") || strings.Contains(item, "}}") {
				return "", ErrMalformedPlan
			}
			lines[i] = fmt.Sprintf("- [P%d] %s", i+1, item)
		}
		return strings.Join(lines, "\n"), nil
	}
	normalizeVerificationItems := func(items []string, minimum, maximum int) (string, error) {
		if len(items) < minimum || len(items) > maximum {
			return "", ErrMalformedPlan
		}
		lines := make([]string, len(items))
		for i, item := range items {
			item = strings.TrimSpace(item)
			if item == "" || strings.ContainsAny(item, "\r\n") || strings.Contains(item, "{{") || strings.Contains(item, "}}") {
				return "", ErrMalformedPlan
			}
			lines[i] = "- " + item
		}
		return strings.Join(lines, "\n"), nil
	}
	steps, err := normalizeStepItems(fields.Steps, 2, 6)
	if err != nil {
		return "", "", "", err
	}
	verification, err := normalizeVerificationItems(fields.Verification, 1, 3)
	if err != nil {
		return "", "", "", err
	}
	return goal, steps, verification, nil
}

func RenderPlan(tmpl string, fields PlanFields, maxOutputBytes int) ([]byte, error) {
	if maxOutputBytes < 1 {
		return nil, ErrMalformedPlan
	}
	if err := ValidatePlanTemplate(tmpl); err != nil {
		return nil, err
	}

	goal, steps, verification, err := normalizePlanFields(fields)
	if err != nil {
		return nil, err
	}
	res := strings.ReplaceAll(tmpl, "{{goal}}", goal)
	res = strings.ReplaceAll(res, "{{steps}}", steps)
	res = strings.ReplaceAll(res, "{{verification}}", verification)
	if strings.Contains(res, "{{") || strings.Contains(res, "}}") || len(res) > maxOutputBytes {
		return nil, ErrMalformedPlan
	}
	return []byte(res), nil
}

func RenderReview(tmpl string, fields ReviewFields, maxOutputBytes int) ([]byte, error) {
	if maxOutputBytes < 1 {
		return nil, ErrMalformedReview
	}
	if err := ValidateReviewTemplate(tmpl); err != nil {
		return nil, err
	}

	itemStatus := strings.TrimSpace(fields.ItemStatus)
	changes := strings.TrimSpace(fields.Changes)
	verification := strings.TrimSpace(fields.Verification)
	deviations := strings.TrimSpace(fields.Deviations)
	res := strings.ReplaceAll(tmpl, "{{item_status}}", itemStatus)
	res = strings.ReplaceAll(res, "{{changes}}", changes)
	res = strings.ReplaceAll(res, "{{verification}}", verification)
	res = strings.ReplaceAll(res, "{{deviations}}", deviations)

	if strings.Contains(res, "{{") || strings.Contains(res, "}}") {
		return nil, ErrMalformedReview
	}

	if len(res) > maxOutputBytes {
		return nil, ErrMalformedReview
	}

	return []byte(res), nil
}