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(`\{\{[^}]*\}\}`)

type ReviewFields struct {
	Checks       string
	Verification string
	Summary      string
}

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

	if strings.Count(tmpl, "{{goal}}") != 1 {
		return fmt.Errorf("%w: must contain {{goal}} exactly once", ErrInvalidTemplate)
	}
	if strings.Count(tmpl, "{{steps}}") != 1 {
		return fmt.Errorf("%w: must contain {{steps}} exactly once", ErrInvalidTemplate)
	}
	if strings.Count(tmpl, "{{verification}}") != 1 {
		return fmt.Errorf("%w: must contain {{verification}} exactly once", ErrInvalidTemplate)
	}

	matches := placeholderRegex.FindAllString(tmpl, -1)
	for _, m := range matches {
		if m != "{{goal}}" && m != "{{steps}}" && m != "{{verification}}" {
			return fmt.Errorf("%w: unknown placeholder %q", ErrInvalidTemplate, m)
		}
	}

	idxGoal := strings.Index(tmpl, "{{goal}}")
	idxSteps := strings.Index(tmpl, "{{steps}}")
	idxVerif := strings.Index(tmpl, "{{verification}}")
	if !(idxGoal < idxSteps && idxSteps < idxVerif) {
		return fmt.Errorf("%w: placeholders must appear in order {{goal}}, {{steps}}, {{verification}}", ErrInvalidTemplate)
	}

	idxPlanH := strings.Index(tmpl, "# Plan")
	idxGoalH := strings.Index(tmpl, "## Goal")
	idxStepsH := strings.Index(tmpl, "## Steps")
	idxVerifH := strings.Index(tmpl, "## Verification")

	if idxPlanH < 0 || idxGoalH < 0 || idxStepsH < 0 || idxVerifH < 0 {
		return fmt.Errorf("%w: missing required headings (# Plan, ## Goal, ## Steps, ## Verification)", ErrInvalidTemplate)
	}
	if !(idxPlanH < idxGoalH && idxGoalH < idxGoal && idxGoal < idxStepsH && idxStepsH < idxSteps && idxSteps < idxVerifH && 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)
	}

	if strings.Count(tmpl, "{{checks}}") != 1 {
		return fmt.Errorf("%w: must contain {{checks}} exactly once", ErrInvalidTemplate)
	}
	if strings.Count(tmpl, "{{verification}}") != 1 {
		return fmt.Errorf("%w: must contain {{verification}} exactly once", ErrInvalidTemplate)
	}
	if strings.Count(tmpl, "{{summary}}") != 1 {
		return fmt.Errorf("%w: must contain {{summary}} exactly once", ErrInvalidTemplate)
	}

	matches := placeholderRegex.FindAllString(tmpl, -1)
	for _, m := range matches {
		if m != "{{checks}}" && m != "{{verification}}" && m != "{{summary}}" {
			return fmt.Errorf("%w: unknown placeholder %q", ErrInvalidTemplate, m)
		}
	}

	idxChecks := strings.Index(tmpl, "{{checks}}")
	idxVerif := strings.Index(tmpl, "{{verification}}")
	idxSumm := strings.Index(tmpl, "{{summary}}")
	if !(idxChecks < idxVerif && idxVerif < idxSumm) {
		return fmt.Errorf("%w: placeholders must appear in order {{checks}}, {{verification}}, {{summary}}", ErrInvalidTemplate)
	}

	idxReviewH := strings.Index(tmpl, "# Review")
	idxResultH := strings.Index(tmpl, "## Result")
	idxPass := strings.Index(tmpl, "PASS")
	idxChecksH := strings.Index(tmpl, "## Checks")
	idxVerifH := strings.Index(tmpl, "## Verification")
	idxSummH := strings.Index(tmpl, "## Summary")

	if idxReviewH < 0 || idxResultH < 0 || idxPass < 0 || idxChecksH < 0 || idxVerifH < 0 || idxSummH < 0 {
		return fmt.Errorf("%w: missing required headings or PASS result", ErrInvalidTemplate)
	}
	if !(idxReviewH < idxResultH && idxResultH < idxPass && idxPass < idxChecksH && idxChecksH < idxChecks && idxChecks < idxVerifH && idxVerifH < idxVerif && idxVerif < idxSummH && 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 {
		if !strings.HasSuffix(rem, f3) {
			return nil, ErrMalformedPlan
		}
		vVerif = rem[:len(rem)-len(f3)]
	}

	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
}