source/pkg/compactcli/run.go

// Package compactcli implements the exact fail-closed whole-package command
// used for compact candidate evaluation. It imports no legacy model or scorer.
package compactcli

import (
	"encoding/json"
	"flag"
	"io"

	"huggingface.co/turenlabs/Vigil/source/pkg/compactonnx"
	"huggingface.co/turenlabs/Vigil/source/pkg/compactscan"
)

const (
	ResultSchema = "vigil.compact-score.v1"
	ErrorSchema  = "vigil.compact-error.v1"
)

type Result struct {
	SchemaVersion        string            `json:"schema_version"`
	InputIndex           int               `json:"input_index"`
	Label                string            `json:"label"`
	MaliciousProbability float64           `json:"malicious_probability"`
	Threshold            float64           `json:"threshold"`
	WholePackage         bool              `json:"whole_package"`
	ModelRequired        bool              `json:"model_required"`
	Package              compactscan.Stats `json:"package"`
	compactonnx.Markers
}

type ErrorResult struct {
	SchemaVersion string `json:"schema_version"`
	ErrorCode     string `json:"error_code"`
	InputIndex    *int   `json:"input_index,omitempty"`
}

// Run executes vigil-compact arguments and returns a process exit code. Errors
// are deliberately aggregate-safe: target paths, member names, and content are
// never emitted.
func Run(arguments []string, stdout, stderr io.Writer) int {
	flags := flag.NewFlagSet("vigil-compact", flag.ContinueOnError)
	flags.SetOutput(io.Discard)
	modelPath := flags.String("model", "", "external compact ONNX candidate")
	metadataPath := flags.String("metadata", "", "adjacent candidate metadata")
	runtimeLibrary := flags.String("runtime-lib", "", "ONNX Runtime shared library")
	requireModel := flags.Bool("require-model", false, "require ONNX model loading (compact mode never falls back)")
	releaseMode := flags.Bool("release", false, "require the embedded released compact model")
	format := flags.String("format", "json", "machine-readable output format (json only)")
	if err := flags.Parse(arguments); err != nil {
		writeError(stderr, "invalid_arguments", nil)
		return 2
	}
	if *format != "json" {
		writeError(stderr, "unsupported_format", nil)
		return 2
	}
	if *releaseMode {
		// No compact model has been promoted. Failing here prevents an external
		// development candidate from being represented as a release artifact.
		writeError(stderr, "released_compact_model_unavailable", nil)
		return 2
	}
	if *modelPath == "" {
		code := "model_required"
		if !*requireModel {
			code = "model_missing_no_fallback"
		}
		writeError(stderr, code, nil)
		return 2
	}
	if len(flags.Args()) == 0 {
		writeError(stderr, "package_required", nil)
		return 2
	}

	model, err := compactonnx.Load(compactonnx.Config{
		ModelPath: *modelPath, MetadataPath: *metadataPath, RuntimeLibrary: *runtimeLibrary,
	})
	if err != nil {
		writeError(stderr, "model_load_failed", nil)
		return 2
	}
	defer model.Close()

	encoder := json.NewEncoder(stdout)
	failed := false
	for index, target := range flags.Args() {
		pkg, stats, err := compactscan.Scan(target)
		if err != nil {
			writeError(stderr, "package_scan_failed", &index)
			failed = true
			continue
		}
		probability, err := model.Score(pkg)
		if err != nil {
			writeError(stderr, "model_score_failed", &index)
			failed = true
			continue
		}
		label := "benign"
		if probability >= model.Threshold() {
			label = "malicious"
		}
		result := Result{
			SchemaVersion: ResultSchema, InputIndex: index, Label: label,
			MaliciousProbability: probability, Threshold: model.Threshold(),
			WholePackage: true, ModelRequired: true, Package: stats,
			Markers: model.Markers(),
		}
		if err := encoder.Encode(result); err != nil {
			writeError(stderr, "output_failed", &index)
			return 1
		}
	}
	if failed {
		return 1
	}
	return 0
}

func writeError(output io.Writer, code string, index *int) {
	if code == "" {
		code = "internal_error"
	}
	_ = json.NewEncoder(output).Encode(ErrorResult{SchemaVersion: ErrorSchema, ErrorCode: code, InputIndex: index})
}