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

## Result
PASS

## Checks
{{checks}}

## Verification
{{verification}}

## Summary
{{summary}}
`

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{"{{checks}}", "{{verification}}", "{{summary}}"}
	reviewLines        = []string{"# Review", "## Result", "PASS", "## Checks", "## Verification", "## Summary"}
)

type ReviewFields struct {
	Checks       string
	Verification string
	Summary      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
	}
	idxChecks, idxVerif, idxSumm := placeholders[0], placeholders[1], placeholders[2]
	if !ascending(idxChecks, idxVerif, idxSumm) {
		return fmt.Errorf("%w: placeholders must appear in order {{checks}}, {{verification}}, {{summary}}", ErrInvalidTemplate)
	}

	lines, err := requireExactLines(tmpl, reviewLines)
	if err != nil {
		return err
	}
	idxReviewH, idxResultH, idxPass := lines[0], lines[1], lines[2]
	idxChecksH, idxVerifH, idxSummH := lines[3], lines[4], lines[5]
	if !ascending(idxReviewH, idxResultH, idxPass, idxChecksH, idxChecks, idxVerifH, idxVerif, idxSummH, idxSumm) {
		return fmt.Errorf("%w: headings and placeholders must follow exact structural order", ErrInvalidTemplate)
	}

	return nil
}

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)]
	}

	trimmedGoal := strings.TrimSpace(vGoal)
	if trimmedGoal == "" || strings.Contains(trimmedGoal, "\n") {
		return nil, ErrMalformedPlan
	}

	trimmedSteps := strings.TrimSpace(vSteps)
	if trimmedSteps == "" {
		return nil, ErrMalformedPlan
	}
	stepLines := strings.Split(trimmedSteps, "\n")
	if len(stepLines) < 2 || len(stepLines) > 6 {
		return nil, ErrMalformedPlan
	}
	for _, l := range stepLines {
		trimmedLine := strings.TrimSpace(l)
		if !strings.HasPrefix(trimmedLine, "- ") || strings.TrimSpace(trimmedLine[2:]) == "" {
			return nil, ErrMalformedPlan
		}
	}

	trimmedVerif := strings.TrimSpace(vVerif)
	if trimmedVerif == "" {
		return nil, ErrMalformedPlan
	}
	verifLines := strings.Split(trimmedVerif, "\n")
	if len(verifLines) < 1 || len(verifLines) > 3 {
		return nil, ErrMalformedPlan
	}
	for _, l := range verifLines {
		trimmedLine := strings.TrimSpace(l)
		if !strings.HasPrefix(trimmedLine, "- ") || strings.TrimSpace(trimmedLine[2:]) == "" {
			return nil, ErrMalformedPlan
		}
	}

	return []byte(rawOutput), 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
	}

	c := strings.TrimSpace(fields.Checks)
	v := strings.TrimSpace(fields.Verification)
	s := strings.TrimSpace(fields.Summary)
	if c == "" || v == "" || s == "" {
		return nil, ErrMalformedReview
	}

	res := strings.ReplaceAll(tmpl, "{{checks}}", c)
	res = strings.ReplaceAll(res, "{{verification}}", v)
	res = strings.ReplaceAll(res, "{{summary}}", s)

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

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

	return []byte(res), nil
}