source/pkg/compactmodel/serialize.go

package compactmodel

import (
	"bytes"
	"errors"
	"fmt"
	"math"
	"sort"
	"strings"
	"unicode/utf8"

	"golang.org/x/text/cases"
	"golang.org/x/text/language"
)

const truncationMarker = "\n[...MIDDLE TRUNCATED...]\n"

// File is one already-scanned package member. Path must be a canonical,
// package-relative slash path. Content is never interpreted as a filesystem
// path and is not mutated.
type File struct {
	Path       string
	Content    []byte
	Executable bool
	// SizeBytes is the original file size when Content is a bounded head/tail
	// sample produced by a trusted filesystem adapter. Zero means Content is
	// complete. A bounded sample is exactly MaxFileBytes bytes: the first and
	// last MaxFileBytes/2 bytes concatenated without a marker.
	SizeBytes int64
}

// Package is a safe, already-scanned whole-package representation. Callers,
// not this package, own archive extraction, symlink handling and IO policy.
type Package struct {
	Files []File
}

// Document is the deterministic serialized text plus the 16 package features.
type Document struct {
	Text       string
	Structured [StructuredFeatures]float32
	BytesRead  int
	FileCount  int
}

var (
	textSuffixes             = stringSet(".bash", ".c", ".cfg", ".conf", ".cpp", ".css", ".csv", ".go", ".h", ".html", ".ini", ".java", ".js", ".json", ".jsx", ".lua", ".md", ".mjs", ".ps1", ".py", ".rb", ".rs", ".sh", ".sql", ".toml", ".ts", ".tsx", ".txt", ".xml", ".yaml", ".yml")
	scriptSuffixes           = stringSet(".bash", ".js", ".mjs", ".ps1", ".py", ".rb", ".sh")
	configNames              = stringSet(".env", ".npmrc", "dockerfile", "makefile", "package.json", "pyproject.toml", "requirements.txt", "settings.json")
	archiveSuffixes          = stringSet(".7z", ".docx", ".gz", ".jar", ".tar", ".war", ".xlsx", ".zip")
	nativeSuffixes           = stringSet(".dll", ".dylib", ".exe", ".node", ".pyc", ".so", ".wasm")
	imageSuffixes            = stringSet(".gif", ".ico", ".jpeg", ".jpg", ".png", ".svg", ".webp")
	structuredConfigSuffixes = stringSet(".cfg", ".conf", ".ini", ".toml", ".yaml", ".yml")
)

func stringSet(values ...string) map[string]struct{} {
	result := make(map[string]struct{}, len(values))
	for _, value := range values {
		result[value] = struct{}{}
	}
	return result
}

func contains(set map[string]struct{}, value string) bool {
	_, ok := set[value]
	return ok
}

func validatePath(value string) error {
	if value == "" {
		return errors.New("path is empty")
	}
	if !utf8.ValidString(value) {
		return errors.New("path is not valid UTF-8")
	}
	if strings.HasPrefix(value, "/") || strings.Contains(value, "\\") {
		return errors.New("path is not a canonical relative slash path")
	}
	for _, r := range value {
		if r < 0x20 || r == 0x7f {
			return errors.New("path contains a control character")
		}
	}
	parts := strings.Split(value, "/")
	for _, part := range parts {
		if part == "" || part == "." || part == ".." {
			return errors.New("path contains an empty or traversal component")
		}
	}
	return nil
}

func pathParts(value string) []string { return strings.Split(value, "/") }

func baseName(value string) string {
	parts := pathParts(value)
	return parts[len(parts)-1]
}

// pythonSuffix matches pathlib.PurePath.suffix for the canonical filenames
// accepted here. In particular, dotfiles and names ending in dots have no
// suffix, unlike path.Ext in Go.
func pythonSuffix(value string) string {
	name := baseName(value)
	if name == "" || strings.HasSuffix(name, ".") {
		return ""
	}
	index := strings.LastIndexByte(name, '.')
	if index <= 0 {
		return ""
	}
	return name[index:]
}

type orderedFile struct {
	file      File
	baseLower string
	suffix    string
	pathLower string
	depth     int
	rank      int
}

func orderFiles(files []File) ([]orderedFile, error) {
	ordered := make([]orderedFile, len(files))
	seen := make(map[string]struct{}, len(files))
	lowerCaser := cases.Lower(language.Und)
	for index, file := range files {
		if err := validatePath(file.Path); err != nil {
			return nil, fmt.Errorf("compactmodel: invalid package path %q: %w", file.Path, err)
		}
		if _, exists := seen[file.Path]; exists {
			return nil, fmt.Errorf("compactmodel: duplicate package path %q", file.Path)
		}
		seen[file.Path] = struct{}{}
		baseLower := lowerCaser.String(baseName(file.Path))
		suffix := lowerCaser.String(pythonSuffix(file.Path))
		rank := 4
		switch {
		case baseLower == "skill.md":
			rank = 0
		case contains(scriptSuffixes, suffix) || contains(configNames, baseLower):
			rank = 1
		case contains(nativeSuffixes, suffix) || contains(archiveSuffixes, suffix):
			rank = 2
		case contains(textSuffixes, suffix):
			rank = 3
		}
		ordered[index] = orderedFile{
			file: file, baseLower: baseLower, suffix: suffix,
			pathLower: lowerCaser.String(file.Path), depth: len(pathParts(file.Path)), rank: rank,
		}
	}
	sort.Slice(ordered, func(left, right int) bool {
		a, b := ordered[left], ordered[right]
		if a.rank != b.rank {
			return a.rank < b.rank
		}
		if a.depth != b.depth {
			return a.depth < b.depth
		}
		if a.pathLower != b.pathLower {
			return a.pathLower < b.pathLower
		}
		return a.file.Path < b.file.Path
	})
	if len(ordered) > MaxPackageFiles {
		ordered = ordered[:MaxPackageFiles]
	}
	return ordered, nil
}

func sampleBytes(data []byte, limit int) []byte {
	if len(data) <= limit {
		return data
	}
	half := limit / 2
	result := make([]byte, 0, half*2+len(truncationMarker))
	result = append(result, data[:half]...)
	result = append(result, truncationMarker...)
	result = append(result, data[len(data)-half:]...)
	return result
}

func sampledContent(file File, limit int) ([]byte, int64, error) {
	size := file.SizeBytes
	if size == 0 {
		size = int64(len(file.Content))
	}
	if size < int64(len(file.Content)) {
		return nil, 0, fmt.Errorf("compactmodel: file %q size %d is smaller than its %d content bytes", file.Path, size, len(file.Content))
	}
	if size == int64(len(file.Content)) {
		return sampleBytes(file.Content, limit), size, nil
	}
	if size <= MaxFileBytes || len(file.Content) != MaxFileBytes {
		return nil, 0, fmt.Errorf("compactmodel: file %q has an invalid bounded sample for size %d", file.Path, size)
	}
	half := limit / 2
	if half > MaxFileBytes/2 {
		return nil, 0, fmt.Errorf("compactmodel: invalid sample limit %d", limit)
	}
	result := make([]byte, 0, half*2+len(truncationMarker))
	result = append(result, file.Content[:half]...)
	result = append(result, truncationMarker...)
	result = append(result, file.Content[len(file.Content)-half:]...)
	return result, size, nil
}

// SelectFiles validates, orders, and applies the contract's package member
// cap without serializing file content. It lets filesystem adapters discover
// every member first and only read the members the frozen contract will use.
func SelectFiles(files []File) ([]File, error) {
	ordered, err := orderFiles(files)
	if err != nil {
		return nil, err
	}
	result := make([]File, len(ordered))
	for index, entry := range ordered {
		result[index] = entry.file
	}
	return result, nil
}

func printableContent(data []byte) (string, bool) {
	if utf8.Valid(data) && !bytes.ContainsRune(data, '\x00') {
		return string(data), false
	}
	var stringsFound []string
	for start := 0; start < len(data); {
		for start < len(data) && (data[start] < 0x20 || data[start] > 0x7e) {
			start++
		}
		end := start
		for end < len(data) && data[end] >= 0x20 && data[end] <= 0x7e {
			end++
		}
		if end-start >= 4 {
			stringsFound = append(stringsFound, string(data[start:end]))
		}
		start = end
	}
	return strings.Join(stringsFound, "\n"), true
}

// Serialize validates and deterministically serializes an already-scanned
// package. It performs no filesystem operations and follows the training
// package-structure-v1 byte accounting exactly.
func Serialize(pkg Package) (Document, error) {
	if err := ValidateRuntime(); err != nil {
		return Document{}, err
	}
	entries, err := orderFiles(pkg.Files)
	if err != nil {
		return Document{}, err
	}

	chunks := make([]string, 0, len(entries)*2)
	totalBytes := 0
	textCount, binaryCount, hiddenCount, executableCount := 0, 0, 0, 0
	scriptCount, configCount, archiveCount := 0, 0, 0
	nativeCount, imageCount, maxDepth := 0, 0, 0

	for _, entry := range entries {
		if totalBytes >= MaxPackageBytes {
			break
		}
		remaining := min(MaxFileBytes, MaxPackageBytes-totalBytes)
		raw, rawSize, err := sampledContent(entry.file, remaining)
		if err != nil {
			return Document{}, err
		}
		totalBytes += len(raw)
		parts := pathParts(entry.file.Path)
		maxDepth = max(maxDepth, len(parts))
		content, binary := printableContent(raw)
		if binary {
			binaryCount++
		} else {
			textCount++
		}
		for _, part := range parts {
			if strings.HasPrefix(part, ".") {
				hiddenCount++
				break
			}
		}
		if entry.file.Executable {
			executableCount++
		}
		if contains(scriptSuffixes, entry.suffix) {
			scriptCount++
		}
		if contains(configNames, entry.baseLower) || contains(structuredConfigSuffixes, entry.suffix) {
			configCount++
		}
		if contains(archiveSuffixes, entry.suffix) {
			archiveCount++
		}
		if contains(nativeSuffixes, entry.suffix) {
			nativeCount++
		}
		if contains(imageSuffixes, entry.suffix) {
			imageCount++
		}
		suffixLabel := entry.suffix
		if suffixLabel == "" {
			suffixLabel = "none"
		}
		chunks = append(chunks,
			fmt.Sprintf("[FILE path=%s suffix=%s bytes=%d binary=%d]", entry.file.Path, suffixLabel, rawSize, boolInt(binary)),
			content,
		)
	}

	fileCount := max(1, textCount+binaryCount)
	suspiciousCount := archiveCount + nativeCount
	structured := [StructuredFeatures]float32{
		float32(math.Min(1, math.Log1p(float64(totalBytes))/math.Log1p(MaxPackageBytes))),
		float32(math.Min(1, float64(fileCount)/float64(MaxPackageFiles))),
		float32(float64(textCount) / float64(fileCount)),
		float32(float64(binaryCount) / float64(fileCount)),
		float32(float64(hiddenCount) / float64(fileCount)),
		float32(float64(executableCount) / float64(fileCount)),
		float32(float64(scriptCount) / float64(fileCount)),
		float32(float64(configCount) / float64(fileCount)),
		float32(float64(archiveCount) / float64(fileCount)),
		float32(float64(nativeCount) / float64(fileCount)),
		float32(float64(imageCount) / float64(fileCount)),
		float32(math.Min(1, float64(maxDepth)/12)),
		float32(float64(suspiciousCount) / float64(fileCount)),
		boolFloat32(fileCount > 1),
		boolFloat32(hasSkillMD(entries)),
		boolFloat32(totalBytes >= MaxPackageBytes),
	}
	return Document{Text: strings.Join(chunks, "\n"), Structured: structured, BytesRead: totalBytes, FileCount: fileCount}, nil
}

func hasSkillMD(entries []orderedFile) bool {
	for _, entry := range entries {
		if entry.baseLower == "skill.md" {
			return true
		}
	}
	return false
}

func boolInt(value bool) int {
	if value {
		return 1
	}
	return 0
}

func boolFloat32(value bool) float32 { return float32(boolInt(value)) }