Agent-System/internal/agent/resultv1.go

package agent

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
)

const (
	ResultV1SchemaVersion = 1
	MaxResultMembers      = 32
	MaxResultMemberBytes  = 2 << 20
	MaxResultBodyBytes    = 8 << 20
	RunnerProtocolLegacy  = "legacy"
	RunnerProtocolTypedV1 = "typed-v1"
)

var ErrTypedResult = fmt.Errorf("agent: typed result_v1 rejected")

var allowedErrorCodes = map[string]struct{}{
	"PROMPT_EMPTY":   {},
	"UNAUTHORIZED":   {},
	"TIMEOUT":        {},
	"CANCELLED":      {},
	"TRANSIENT":      {},
	"MISSING_TARGET": {},
	"RUN_FAILED":     {},
}

var retryableErrorCodes = map[string]struct{}{
	"TIMEOUT":   {},
	"CANCELLED": {},
	"TRANSIENT": {},
}

type ResultV1Member struct {
	Kind   string `json:"kind"`
	Path   string `json:"path"`
	Digest string `json:"digest"`
	Size   int    `json:"size"`
}

type ResultV1 struct {
	SchemaVersion int              `json:"schema_version"`
	RunID         string           `json:"run_id"`
	Outcome       string           `json:"outcome"`
	Retryable     bool             `json:"retryable"`
	ErrorCode     string           `json:"error_code,omitempty"`
	Members       []ResultV1Member `json:"members"`
	ContentDigest string           `json:"content_digest"`
}

func SafeRelPath(p string) bool {
	p = strings.TrimSpace(p)
	if p == "" || filepath.IsAbs(p) {
		return false
	}
	clean := filepath.ToSlash(filepath.Clean(p))
	if clean == "." || clean == ".." || strings.HasPrefix(clean, "../") || strings.Contains(p, "..") {
		return false
	}
	return true
}

func DigestBytes(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

func (r ResultV1) MembersDigest() string {
	type row struct {
		Kind   string `json:"kind"`
		Path   string `json:"path"`
		Digest string `json:"digest"`
		Size   int    `json:"size"`
	}
	members := r.Members
	if members == nil {
		members = []ResultV1Member{}
	}
	rows := make([]row, 0, len(members))
	for _, m := range members {
		rows = append(rows, row{m.Kind, m.Path, m.Digest, m.Size})
	}
	raw, _ := json.Marshal(rows)
	return DigestBytes(raw)
}

func (r ResultV1) Validate(expectedRunID string, patches map[string][]byte) error {
	if r.SchemaVersion != ResultV1SchemaVersion {
		return fmt.Errorf("%w: schema", ErrTypedResult)
	}
	if strings.TrimSpace(r.RunID) == "" {
		return fmt.Errorf("%w: run_id", ErrTypedResult)
	}
	if expectedRunID != "" && r.RunID != expectedRunID {
		return fmt.Errorf("%w: run mismatch", ErrTypedResult)
	}
	if r.Outcome != "success" && r.Outcome != "error" {
		return fmt.Errorf("%w: outcome", ErrTypedResult)
	}
	if r.Outcome == "success" && strings.TrimSpace(r.ErrorCode) != "" {
		return fmt.Errorf("%w: error_code", ErrTypedResult)
	}
	if r.Outcome == "error" {
		if _, ok := allowedErrorCodes[r.ErrorCode]; !ok {
			return fmt.Errorf("%w: error_code", ErrTypedResult)
		}
	}
	if len(r.Members) > MaxResultMembers {
		return fmt.Errorf("%w: members", ErrTypedResult)
	}
	total := 0
	seen := map[string]struct{}{}
	for _, m := range r.Members {
		if m.Kind != "file" || !SafeRelPath(m.Path) {
			return fmt.Errorf("%w: member path", ErrTypedResult)
		}
		if _, dup := seen[m.Path]; dup {
			return fmt.Errorf("%w: duplicate member", ErrTypedResult)
		}
		seen[m.Path] = struct{}{}
		if m.Size < 0 || m.Size > MaxResultMemberBytes {
			return fmt.Errorf("%w: member size", ErrTypedResult)
		}
		total += m.Size
		body, ok := patches[m.Path]
		if !ok {
			return fmt.Errorf("%w: member missing body", ErrTypedResult)
		}
		if len(body) != m.Size || DigestBytes(body) != m.Digest {
			return fmt.Errorf("%w: digest", ErrTypedResult)
		}
	}
	if total > MaxResultBodyBytes {
		return fmt.Errorf("%w: body", ErrTypedResult)
	}
	if r.ContentDigest != r.MembersDigest() {
		return fmt.Errorf("%w: content_digest", ErrTypedResult)
	}
	return nil
}

func TypedRetryable(r ResultV1) bool {
	if r.Outcome != "error" || !r.Retryable {
		return false
	}
	_, ok := retryableErrorCodes[r.ErrorCode]
	return ok
}

func BuildResultV1(runID, outcome string, retryable bool, errorCode string, patches []struct {
	RelPath string
	Content []byte
}) ResultV1 {
	members := make([]ResultV1Member, 0, len(patches))
	for _, p := range patches {
		if !SafeRelPath(p.RelPath) {
			continue
		}
		members = append(members, ResultV1Member{
			Kind: "file", Path: p.RelPath, Digest: DigestBytes(p.Content), Size: len(p.Content),
		})
	}
	r := ResultV1{
		SchemaVersion: ResultV1SchemaVersion,
		RunID:         runID,
		Outcome:       outcome,
		Retryable:     retryable,
		ErrorCode:     errorCode,
		Members:       members,
	}
	r.ContentDigest = r.MembersDigest()
	return r
}