iop/packages/go/singlerequesttemplate/template.go

package singlerequesttemplate

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

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 ParsePlan(tmpl string, rawOutput string, maxOutputBytes int) ([]byte, error) {
	if maxOutputBytes < 1 || len(rawOutput) > maxOutputBytes {
		return nil, ErrMalformedPlan
	}
	if err := ValidatePlanTemplate(tmpl); err != nil {
		return nil, err
	}

	if strings.Contains(rawOutput, "{{") || strings.Contains(rawOutput, "}}") {
		return nil, ErrMalformedPlan
	}

	idxGoalPlaceholder := strings.Index(tmpl, "{{goal}}")
	idxStepsPlaceholder := strings.Index(tmpl, "{{steps}}")
	idxVerifPlaceholder := strings.Index(tmpl, "{{verification}}")

	f0 := tmpl[:idxGoalPlaceholder]
	f1 := tmpl[idxGoalPlaceholder+len("{{goal}}") : idxStepsPlaceholder]
	f2 := tmpl[idxStepsPlaceholder+len("{{steps}}") : idxVerifPlaceholder]
	f3 := tmpl[idxVerifPlaceholder+len("{{verification}}"):]

	if !strings.HasPrefix(rawOutput, f0) {
		return nil, ErrMalformedPlan
	}
	rem := rawOutput[len(f0):]

	i1 := strings.Index(rem, f1)
	if i1 < 0 {
		return nil, ErrMalformedPlan
	}
	vGoal := rem[:i1]
	rem = rem[i1+len(f1):]

	i2 := strings.Index(rem, f2)
	if i2 < 0 {
		return nil, ErrMalformedPlan
	}
	vSteps := rem[:i2]
	rem = rem[i2+len(f2):]

	var vVerif string
	if f3 == "" {
		vVerif = rem
	} else {
		suffix := f3
		if !strings.HasSuffix(rem, suffix) {
			// Provider chat APIs commonly omit the model's final line feed. Treat
			// only that last byte as optional; all other static suffix text must
			// still match the configured template exactly.
			if !strings.HasSuffix(f3, "\n") {
				return nil, ErrMalformedPlan
			}
			suffix = strings.TrimSuffix(f3, "\n")
			if !strings.HasSuffix(rem, suffix) {
				return nil, ErrMalformedPlan
			}
		}
		vVerif = rem[:len(rem)-len(suffix)]
	}

	if _, _, _, err := normalizePlanSections(vGoal, vSteps, vVerif); err != nil {
		return nil, err
	}

	return []byte(rawOutput), nil
}

func normalizePlanSections(goal, steps, verification string) (string, string, string, error) {
	goal = strings.TrimSpace(goal)
	steps = strings.TrimSpace(steps)
	verification = strings.TrimSpace(verification)
	if goal == "" || strings.ContainsAny(goal, "\r\n") || steps == "" || verification == "" {
		return "", "", "", ErrMalformedPlan
	}
	if strings.Contains(goal, "{{") || strings.Contains(goal, "}}") || strings.Contains(steps, "{{") || strings.Contains(steps, "}}") || strings.Contains(verification, "{{") || strings.Contains(verification, "}}") {
		return "", "", "", ErrMalformedPlan
	}

	normalizeStepBullets := func(value string, minimum, maximum int) (string, error) {
		lines := strings.Split(value, "\n")
		if len(lines) < minimum || len(lines) > maximum {
			return "", ErrMalformedPlan
		}
		for i, line := range lines {
			line = strings.TrimSpace(line)
			prefix := fmt.Sprintf("- [P%d] ", i+1)
			if !strings.HasPrefix(line, prefix) || strings.TrimSpace(line[len(prefix):]) == "" {
				return "", ErrMalformedPlan
			}
			lines[i] = line
		}
		return strings.Join(lines, "\n"), nil
	}

	normalizeVerificationBullets := func(value string, minimum, maximum int) (string, error) {
		lines := strings.Split(value, "\n")
		if len(lines) < minimum || len(lines) > maximum {
			return "", ErrMalformedPlan
		}
		for i, line := range lines {
			line = strings.TrimSpace(line)
			if !strings.HasPrefix(line, "- ") || strings.TrimSpace(line[2:]) == "" {
				return "", ErrMalformedPlan
			}
			lines[i] = line
		}
		return strings.Join(lines, "\n"), nil
	}

	steps, err := normalizeStepBullets(steps, 2, 6)
	if err != nil {
		return "", "", "", err
	}
	verification, err = normalizeVerificationBullets(verification, 1, 3)
	if err != nil {
		return "", "", "", err
	}
	return goal, steps, verification, nil
}

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)
	if itemStatus == "" || changes == "" || verification == "" || deviations == "" {
		return nil, ErrMalformedReview
	}

	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
}

var planStepIDRegex = regexp.MustCompile(`(?m)^- \[P(\d+)\]`)

// PlanItemIDs extracts the deterministic P1..Pn step IDs from a rendered Plan
// document. IDs must start at P1 and increment with no gaps, duplicates, or
// out-of-order entries.
func PlanItemIDs(plan []byte) ([]string, error) {
	if len(plan) == 0 || !utf8.Valid(plan) {
		return nil, ErrMalformedPlan
	}
	text := string(plan)
	for _, heading := range planHeadings {
		if _, count := exactLineOffsets(text, heading); count != 1 {
			return nil, ErrMalformedPlan
		}
	}
	goalStart := strings.Index(text, "## Goal") + len("## Goal")
	stepsHeading := strings.Index(text, "## Steps")
	stepsStart := stepsHeading + len("## Steps")
	verificationHeading := strings.Index(text, "## Verification")
	if goalStart < len("## Goal") || stepsHeading < 0 || verificationHeading < 0 || goalStart >= stepsHeading || stepsStart >= verificationHeading || strings.TrimSpace(text[goalStart:stepsHeading]) == "" || strings.TrimSpace(text[verificationHeading+len("## Verification"):]) == "" {
		return nil, ErrMalformedPlan
	}
	steps := strings.TrimSpace(text[stepsStart:verificationHeading])
	lines := strings.Split(steps, "\n")
	if len(lines) < 2 || len(lines) > 6 {
		return nil, ErrMalformedPlan
	}
	matches := planStepIDRegex.FindAllStringSubmatch(steps, -1)
	if len(matches) != len(lines) {
		return nil, ErrMalformedPlan
	}
	ids := make([]string, 0, len(matches))
	for i, match := range matches {
		expected := fmt.Sprintf("P%d", i+1)
		actual := "P" + match[1]
		line := strings.TrimSpace(lines[i])
		prefix := fmt.Sprintf("- [%s] ", expected)
		if actual != expected || !strings.HasPrefix(line, prefix) || strings.TrimSpace(line[len(prefix):]) == "" {
			return nil, ErrMalformedPlan
		}
		ids = append(ids, expected)
	}
	return ids, nil
}

var reviewItemLineRegex = regexp.MustCompile(`(?m)^- (P\d+): (.+)$`)

// ValidateReviewHandoff validates a rendered REVIEW handoff document against
// the supplied plan IDs. Every plan ID must appear exactly once in the Worker
// Item Status section with status "completed", and no unknown or duplicate
// IDs are permitted. All four required sections must be present with
// non-empty content. If the PLAN has no deviations, the Deviations section
// must still contain an explicit entry (conventionally "None").
//
// The Worker Item Status section is validated by exact line inventory rather
// than regex filtering: the entire section (excluding the heading) is split on
// newlines, every resulting line must be non-empty, and each line must match
// its corresponding plan ID in the form "- Pn: completed". This rejects prose
// injected between status lines, blank lines, malformed bullets, and any
// out-of-order or duplicate entries in a single pass.
func ValidateReviewHandoff(content []byte, planIDs []string) error {
	if len(content) == 0 || !utf8.Valid(content) || len(planIDs) == 0 {
		return ErrMalformedReview
	}
	text := string(content)

	requiredSections := []string{"# Review", "## Worker Item Status", "## Worker Changes", "## Worker Verification", "## Deviations"}
	for _, section := range requiredSections {
		if _, count := exactLineOffsets(text, section); count != 1 {
			return ErrMalformedReview
		}
	}
	if err := rejectUnknownMarkdownHeadings(text, reviewHeadings); err != nil {
		return ErrMalformedReview
	}
	sectionContent := func(heading, next string) string {
		start := strings.Index(text, heading) + len(heading)
		end := len(text)
		if next != "" {
			if index := strings.Index(text[start:], next); index >= 0 {
				end = start + index
			}
		}
		return strings.TrimSpace(text[start:end])
	}
	if sectionContent("## Worker Item Status", "\n## Worker Changes") == "" ||
		sectionContent("## Worker Changes", "\n## Worker Verification") == "" ||
		sectionContent("## Worker Verification", "\n## Deviations") == "" ||
		sectionContent("## Deviations", "") == "" {
		return ErrMalformedReview
	}

	statusHeading := "## Worker Item Status"
	statusIdx := strings.Index(text, statusHeading)
	if statusIdx < 0 {
		return ErrMalformedReview
	}
	statusSection := sectionContent(statusHeading, "\n## Worker Changes")

	// Exact line inventory: every line in the status section must correspond
	// to one plan ID in order, with the grammar "- Pn: completed". Blank lines,
	// prose, malformed bullets, and out-of-order or duplicate entries are all
	// rejected because the line count and each line's content are compared
	// directly against the plan ID inventory.
	lines := strings.Split(statusSection, "\n")
	if len(lines) != len(planIDs) {
		return ErrMalformedReview
	}
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			return ErrMalformedReview
		}
		expected := fmt.Sprintf("- %s: completed", planIDs[i])
		if trimmed != expected {
			return ErrMalformedReview
		}
	}

	return nil
}