ThothII/tools/thothctl/internal/workspaceops/operations.go

// Package workspaceops defines the deliberately closed host workspace contract.
package workspaceops

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"

	"github.com/aritmolab/thothii/tools/thothctl/internal/compose"
	"github.com/aritmolab/thothii/tools/thothctl/internal/config"
	"github.com/aritmolab/thothii/tools/thothctl/internal/safeio"
)

const CodeRegistryBootstrapRecoveryConflict = "registry_bootstrap_recovery_conflict"
const (
	maxSQLFile         = 1 << 20
	maxSQLTotal        = 16 << 20
	maxAssumptions     = 256
	maxAssumptionBytes = 256
	maxResult          = 1 << 20
	maxCandidate       = 700 << 10
)

var workspaceIDPattern = regexp.MustCompile(`^[a-z][a-z0-9-]{2,62}$`)
var runIDPattern = regexp.MustCompile(`^[0-9a-f]{32}$`)
var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)

type Command interface{ workspaceCommand() }
type InspectCommand struct {
	WorkspaceID string
	JSON        bool
}
type DwhRequest struct {
	WorkspaceID string
	Resume      string
	JSON        bool
}
type SuggestFksRequest struct {
	WorkspaceID string
	FromSQL     []string
	Assume      []string
	Output      string
	JSON        bool
}
type CheckSchemaRequest struct {
	WorkspaceID        string
	Resume             string
	Annotations        string
	ReviewedCandidates string
	JSON               bool
}
type IndexSchemaRequest struct {
	WorkspaceID string
	JSON        bool
}
type EvidenceRequest struct {
	WorkspaceID string
	DryRun      bool
	Resume      string
	JSON        bool
}
type RunRequest struct {
	WorkspaceID string
	Resume      string
	JSON        bool
}

func (InspectCommand) workspaceCommand()     {}
func (DwhRequest) workspaceCommand()         {}
func (SuggestFksRequest) workspaceCommand()  {}
func (CheckSchemaRequest) workspaceCommand() {}
func (IndexSchemaRequest) workspaceCommand() {}
func (EvidenceRequest) workspaceCommand()    {}
func (RunRequest) workspaceCommand()         {}

type Result struct {
	SchemaVersion      int                `json:"schemaVersion"`
	Status             string             `json:"status"`
	Code               string             `json:"code"`
	WorkspaceID        string             `json:"workspaceId"`
	WorkspaceRevision  string             `json:"workspaceRevision"`
	DescriptorBlob     string             `json:"descriptorBlob"`
	Operation          string             `json:"operation"`
	RunID              string             `json:"runId,omitempty"`
	ChildRuns          map[string]string  `json:"childRuns,omitempty"`
	CompletedStages    []string           `json:"completedStages"`
	Counts             map[string]int     `json:"counts,omitempty"`
	ArtifactIdentities []ArtifactIdentity `json:"artifactIdentities,omitempty"`
	Warnings           []string           `json:"warnings,omitempty"`
}
type ArtifactIdentity struct {
	Kind   string `json:"kind"`
	Digest string `json:"digest"`
}

func invalid(msg string) (Command, error) { return nil, errors.New(msg) }
func requireWorkspace(v string) error {
	if !workspaceIDPattern.MatchString(v) {
		return errors.New("--workspace must be a valid workspace ID")
	}
	return nil
}
func requireRun(v string) error {
	if !runIDPattern.MatchString(v) {
		return errors.New("--resume must be a lowercase 32-hex run ID")
	}
	return nil
}
func one(args []string, i *int, flag string) (string, error) {
	if *i+1 >= len(args) || strings.HasPrefix(args[*i+1], "--") {
		return "", fmt.Errorf("%s requires a value", flag)
	}
	*i = *i + 1
	return args[*i], nil
}
func ParseWorkspaceCommand(args []string) (Command, error) {
	if len(args) < 2 || args[0] != "workspace" {
		return invalid("workspace command is required")
	}
	area, action := args[1], ""
	rest := args[2:]
	if area == "preprocess" || area == "schema" {
		if len(rest) == 0 {
			return invalid("workspace group requires an operation")
		}
		action = rest[0]
		rest = rest[1:]
	} else {
		action = area
	}
	if action != "inspect" && action != "dwh" && action != "suggest-fks" && action != "check" && action != "index-schema" && action != "evidence" && action != "run" {
		return invalid("unknown workspace command")
	}
	var ws, resume, annotations, reviewed, output string
	var jsonOut, dry bool
	var sql, assume []string
	seen := map[string]bool{}
	for i := 0; i < len(rest); i++ {
		f := rest[i]
		if f == "--json" {
			if seen[f] {
				return invalid("duplicate --json")
			}
			seen[f] = true
			jsonOut = true
			continue
		}
		switch f {
		case "--workspace":
			if seen[f] {
				return invalid("duplicate --workspace")
			}
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			ws = v
		case "--resume":
			if seen[f] {
				return invalid("duplicate --resume")
			}
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			resume = v
		case "--from-sql":
			if action != "suggest-fks" {
				return invalid("--from-sql is valid only for schema suggest-fks")
			}
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			if len(sql) >= 32 {
				return invalid("too many --from-sql files")
			}
			sql = append(sql, v)
		case "--assume":
			if action != "suggest-fks" {
				return invalid("--assume is valid only for schema suggest-fks")
			}
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			if len(assume) >= maxAssumptions || len([]byte(v)) > maxAssumptionBytes {
				return invalid("--assume exceeds limit")
			}
			assume = append(assume, v)
		case "--output":
			if action != "suggest-fks" {
				return invalid("--output is valid only for schema suggest-fks")
			}
			if seen[f] {
				return invalid("duplicate --output")
			}
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			output = v
		case "--annotations":
			if action != "check" {
				return invalid("--annotations is valid only for schema check")
			}
			if seen[f] {
				return invalid("duplicate --annotations")
			}
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			annotations = v
		case "--reviewed-candidates":
			if action != "check" {
				return invalid("--reviewed-candidates is valid only for schema check")
			}
			if seen[f] {
				return invalid("duplicate --reviewed-candidates")
			}
			seen[f] = true
			v, e := one(rest, &i, f)
			if e != nil {
				return nil, e
			}
			reviewed = v
		case "--dry-run":
			if action != "evidence" {
				return invalid("--dry-run is valid only for evidence")
			}
			if seen[f] {
				return invalid("duplicate --dry-run")
			}
			seen[f] = true
			dry = true
		default:
			return invalid("unknown workspace option")
		}
	}
	if err := requireWorkspace(ws); err != nil {
		return nil, err
	}
	if action != "inspect" && resume != "" {
		if err := requireRun(resume); err != nil {
			return nil, err
		}
	}
	if resume != "" && action != "dwh" && action != "evidence" && action != "run" && action != "check" {
		return invalid("--resume is not valid for this workspace command")
	}
	if action == "inspect" && resume != "" {
		return invalid("inspect does not accept --resume")
	}
	if action == "check" {
		if resume == "" {
			return invalid("schema check requires --resume")
		}
		if (annotations == "") != (reviewed == "") {
			return invalid("--annotations and --reviewed-candidates must be supplied together")
		}
		if reviewed != "" && !digestPattern.MatchString(reviewed) {
			return invalid("--reviewed-candidates must be sha256:<hex>")
		}
	}
	if action != "suggest-fks" && (len(sql) > 0 || len(assume) > 0 || output != "") {
		return invalid("schema options are valid only for suggest-fks")
	}
	switch action {
	case "inspect":
		return InspectCommand{ws, jsonOut}, nil
	case "dwh":
		return DwhRequest{ws, resume, jsonOut}, nil
	case "suggest-fks":
		return SuggestFksRequest{ws, sql, assume, output, jsonOut}, nil
	case "check":
		return CheckSchemaRequest{ws, resume, annotations, reviewed, jsonOut}, nil
	case "index-schema":
		return IndexSchemaRequest{ws, jsonOut}, nil
	case "evidence":
		return EvidenceRequest{ws, dry, resume, jsonOut}, nil
	default:
		return RunRequest{ws, resume, jsonOut}, nil
	}
}

type inputEnvelope struct {
	SchemaVersion      int        `json:"schemaVersion"`
	Operation          string     `json:"operation"`
	WorkspaceID        string     `json:"workspaceId"`
	Resume             string     `json:"resume,omitempty"`
	SQL                []sqlInput `json:"sql,omitempty"`
	Assume             []string   `json:"assume,omitempty"`
	Annotations        string     `json:"annotations,omitempty"`
	ReviewedCandidates string     `json:"reviewedCandidates,omitempty"`
	DryRun             bool       `json:"dryRun,omitempty"`
}
type sqlInput struct {
	Name    string `json:"name"`
	Content string `json:"content"`
}

func operationName(c Command) string {
	switch c.(type) {
	case InspectCommand:
		return "inspect"
	case DwhRequest:
		return "dwh"
	case SuggestFksRequest:
		return "suggest-fks"
	case CheckSchemaRequest:
		return "check"
	case IndexSchemaRequest:
		return "index-schema"
	case EvidenceRequest:
		return "evidence"
	default:
		return "run"
	}
}
func Run(ctx context.Context, installation config.Installation, runner compose.Runner, command Command, stdin io.Reader) (Result, error) {
	env := inputEnvelope{SchemaVersion: 1, Operation: operationName(command)}
	var outputPath string
	switch c := command.(type) {
	case InspectCommand:
		env.WorkspaceID = c.WorkspaceID
	case DwhRequest:
		env.WorkspaceID, env.Resume = c.WorkspaceID, c.Resume
	case SuggestFksRequest:
		env.WorkspaceID, env.Assume, outputPath = c.WorkspaceID, c.Assume, c.Output
		if outputPath != "" {
			if err := safeio.ValidateCanonicalOutputPath(outputPath); err != nil {
				return Result{}, errors.New("unsafe output file")
			}
		}
		var total int64
		for _, path := range c.FromSQL {
			b, e := safeio.ReadCanonicalUTF8(path, maxSQLFile)
			if e != nil {
				return Result{}, errors.New("unsafe SQL input")
			}
			total += int64(len(b))
			if total > maxSQLTotal {
				return Result{}, errors.New("SQL input exceeds limit")
			}
			env.SQL = append(env.SQL, sqlInput{path, string(b)})
		}
		if len(env.SQL) > 0 {
			var n int
			for _, x := range env.SQL {
				n += len(x.Content)
			}
			if n > maxCandidate {
				return Result{}, errors.New("candidate input exceeds limit")
			}
		}
	case CheckSchemaRequest:
		env.WorkspaceID, env.Resume, env.Annotations, env.ReviewedCandidates = c.WorkspaceID, c.Resume, c.Annotations, c.ReviewedCandidates
		if c.Annotations != "" {
			b, e := safeio.ReadCanonicalUTF8(c.Annotations, 16<<20)
			if e != nil {
				return Result{}, errors.New("unsafe annotation input")
			}
			env.Annotations = string(b)
		}
	case IndexSchemaRequest:
		env.WorkspaceID = c.WorkspaceID
	case EvidenceRequest:
		env.WorkspaceID, env.Resume, env.DryRun = c.WorkspaceID, c.Resume, c.DryRun
	case RunRequest:
		env.WorkspaceID, env.Resume = c.WorkspaceID, c.Resume
	default:
		return Result{}, errors.New("unsupported workspace command")
	}
	payload, e := json.Marshal(env)
	if e != nil {
		return Result{}, e
	}
	if stdin != nil {
		payload, e = io.ReadAll(io.LimitReader(stdin, 1<<20+1))
		if e != nil {
			return Result{}, e
		}
		if len(payload) > 1<<20 {
			return Result{}, errors.New("request exceeds limit")
		}
		var supplied inputEnvelope
		d := json.NewDecoder(bytes.NewReader(payload))
		d.DisallowUnknownFields()
		if e = d.Decode(&supplied); e != nil || supplied.SchemaVersion != 1 || supplied.Operation != env.Operation || supplied.WorkspaceID != env.WorkspaceID {
			return Result{}, errors.New("invalid workspace request")
		}
	}

	args := installation.ComposeArgs("run", "--rm", "--no-deps", "workspace-maintenance", "--operation", operationName(command), "--workspace", env.WorkspaceID)
	cr, runErr := runner.RunBounded(ctx, args, bytes.NewReader(payload), compose.CaptureLimits{StdoutBytes: 1 << 20, StderrBytes: 64 << 10})
	if runErr != nil && (errors.Is(runErr, compose.ErrOutputLimit) || cr.Stdout == "") {
		return Result{}, runErr
	}
	if len(cr.Stdout) > maxResult {
		return Result{}, errors.New("invalid workspace result")
	}
	var result Result
	dec := json.NewDecoder(strings.NewReader(cr.Stdout))
	dec.DisallowUnknownFields()
	if e := dec.Decode(&result); e != nil {
		return Result{}, errors.New("invalid workspace result")
	}
	var extra any
	if e := dec.Decode(&extra); e != io.EOF {
		return Result{}, errors.New("invalid workspace result")
	}
	if result.SchemaVersion != 1 || result.WorkspaceID != env.WorkspaceID || result.Operation != operationName(command) || !validStatus(result.Status) || !validCode(result.Code) {
		return Result{}, errors.New("invalid workspace result")
	}
	if outputPath != "" {
		if len(cr.Stdout) > maxCandidate {
			return Result{}, errors.New("candidate export exceeds limit")
		}
		if err := safeio.WriteCanonicalExclusive(outputPath, []byte(cr.Stdout), 0o600); err != nil {
			return Result{}, errors.New("unsafe output file")
		}
	}

	return result, nil
}
func DigestBytes(b []byte) string { s := sha256.Sum256(b); return "sha256:" + hex.EncodeToString(s[:]) }

func validStatus(v string) bool {
	switch v {
	case "succeeded", "unchanged", "dry_run", "blocked", "failed":
		return true
	}
	return false
}
func validCode(v string) bool {
	switch v {
	case "ok", "workspace_not_found", "workspace_not_activatable", "binding_missing", "preprocessing_conflict", "preprocessing_resume_mismatch", "manual_review_required", "evidence_materialization_required", "effective_config_mismatch", "semantic_index_incompatible", "annotation_invalid", "egress_policy_refused", CodeRegistryBootstrapRecoveryConflict:
		return true
	}
	return false
}