weatherreporter/internal/adapters/scriptorium/runner.go

// Package scriptorium adapts the external scriptorium CLI.
package scriptorium

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

const maxCapturedOutputBytes = 1024 * 1024

type CommandRunner interface {
	Run(ctx context.Context, name string, args []string, timeout time.Duration) (CommandResult, error)
}

type CommandResult struct {
	Stdout          []byte
	Stderr          []byte
	StdoutTruncated bool
	StderrTruncated bool
	ExitCode        int
}

type ExecRunner struct{}

func (ExecRunner) Run(ctx context.Context, name string, args []string, timeout time.Duration) (CommandResult, error) {
	runCtx := ctx
	cancel := func() {}
	if timeout > 0 {
		runCtx, cancel = context.WithTimeout(ctx, timeout)
	}
	defer cancel()

	cmd := exec.CommandContext(runCtx, name, args...)
	stdout := &limitedBuffer{limit: maxCapturedOutputBytes}
	stderr := &limitedBuffer{limit: maxCapturedOutputBytes}
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	err := cmd.Run()
	result := CommandResult{
		Stdout:          stdout.Bytes(),
		Stderr:          stderr.Bytes(),
		StdoutTruncated: stdout.Truncated(),
		StderrTruncated: stderr.Truncated(),
		ExitCode:        0,
	}
	if err == nil {
		return result, nil
	}
	if runCtx.Err() != nil {
		return result, runCtx.Err()
	}
	if exitErr, ok := err.(*exec.ExitError); ok {
		result.ExitCode = exitErr.ExitCode()
		return result, nil
	}
	return result, err
}

type Runner struct {
	Binary     string
	ConfigPath string
	Profile    string
	Timeout    time.Duration
	ExtraArgs  []string
	Commands   CommandRunner
}

type RenderRequest struct {
	PromptID        string
	DataPackagePath string
}

type RunRequest struct {
	PromptID        string
	DataPackagePath string
	OutputPath      string
}

type RenderResult struct {
	Command         []string `json:"command"`
	Stdout          string   `json:"stdout"`
	Stderr          string   `json:"stderr"`
	StdoutTruncated bool     `json:"stdoutTruncated,omitempty"`
	StderrTruncated bool     `json:"stderrTruncated,omitempty"`
	ExitCode        int      `json:"exitCode"`
}

type RunResult struct {
	Command         []string `json:"command"`
	Stdout          string   `json:"stdout"`
	Stderr          string   `json:"stderr"`
	StdoutTruncated bool     `json:"stdoutTruncated,omitempty"`
	StderrTruncated bool     `json:"stderrTruncated,omitempty"`
	ExitCode        int      `json:"exitCode"`
	OutputPath      string   `json:"outputPath"`
}

func (r Runner) Render(ctx context.Context, req RenderRequest) (*RenderResult, error) {
	if req.PromptID == "" {
		return nil, fmt.Errorf("prompt id is required")
	}
	if req.DataPackagePath == "" {
		return nil, fmt.Errorf("data package path is required")
	}
	binary := r.Binary
	if binary == "" {
		binary = "scriptorium"
	}
	commands := r.Commands
	if commands == nil {
		commands = ExecRunner{}
	}
	args := r.renderArgs(req)
	commandResult, err := commands.Run(ctx, binary, args, r.Timeout)
	if err != nil {
		return nil, fmt.Errorf("run scriptorium render: %w", err)
	}
	result := &RenderResult{
		Command:         append([]string{binary}, args...),
		Stdout:          string(commandResult.Stdout),
		Stderr:          string(commandResult.Stderr),
		StdoutTruncated: commandResult.StdoutTruncated,
		StderrTruncated: commandResult.StderrTruncated,
		ExitCode:        commandResult.ExitCode,
	}
	if commandResult.ExitCode != 0 {
		return result, fmt.Errorf("scriptorium render exited with code %d: %s", commandResult.ExitCode, result.Stderr)
	}
	return result, nil
}

func (r Runner) Run(ctx context.Context, req RunRequest) (*RunResult, error) {
	if req.PromptID == "" {
		return nil, fmt.Errorf("prompt id is required")
	}
	if req.DataPackagePath == "" {
		return nil, fmt.Errorf("data package path is required")
	}
	if req.OutputPath == "" {
		return nil, fmt.Errorf("output path is required")
	}
	binary := r.Binary
	if binary == "" {
		binary = "scriptorium"
	}
	commands := r.Commands
	if commands == nil {
		commands = ExecRunner{}
	}
	args := r.runArgs(req)
	commandResult, err := commands.Run(ctx, binary, args, r.Timeout)
	if err != nil {
		return nil, fmt.Errorf("run scriptorium: %w", err)
	}
	result := &RunResult{
		Command:         append([]string{binary}, args...),
		Stdout:          string(commandResult.Stdout),
		Stderr:          string(commandResult.Stderr),
		StdoutTruncated: commandResult.StdoutTruncated,
		StderrTruncated: commandResult.StderrTruncated,
		ExitCode:        commandResult.ExitCode,
		OutputPath:      req.OutputPath,
	}
	if commandResult.ExitCode != 0 {
		return result, fmt.Errorf("scriptorium run exited with code %d: %s", commandResult.ExitCode, result.Stderr)
	}
	return result, nil
}

func (r Runner) renderArgs(req RenderRequest) []string {
	args := []string{"render"}
	if r.ConfigPath != "" {
		args = append(args, "--config", r.ConfigPath)
	}
	if r.Profile != "" {
		args = append(args, "--profile", r.Profile)
	}
	args = append(args,
		"--prompt", req.PromptID,
		"--input", "data_package="+req.DataPackagePath,
		"--format", "json",
	)
	args = append(args, r.ExtraArgs...)
	return args
}

func (r Runner) runArgs(req RunRequest) []string {
	args := []string{"run"}
	if r.ConfigPath != "" {
		args = append(args, "--config", r.ConfigPath)
	}
	if r.Profile != "" {
		args = append(args, "--profile", r.Profile)
	}
	args = append(args,
		"--prompt", req.PromptID,
		"--input", "data_package="+req.DataPackagePath,
		"--out", req.OutputPath,
	)
	args = append(args, r.ExtraArgs...)
	return args
}

func SaveRenderResult(path string, result *RenderResult) error {
	if result == nil {
		return fmt.Errorf("render result is required")
	}
	data, err := json.MarshalIndent(result, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal render result: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return fmt.Errorf("create preflight directory %q: %w", filepath.Dir(path), err)
	}
	tmp, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*.tmp")
	if err != nil {
		return fmt.Errorf("create temporary preflight file: %w", err)
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)

	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return fmt.Errorf("write temporary preflight file: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("close temporary preflight file: %w", err)
	}
	if err := os.Rename(tmpName, path); err != nil {
		return fmt.Errorf("save preflight %q: %w", path, err)
	}
	return nil
}

type limitedBuffer struct {
	data      []byte
	limit     int
	truncated bool
}

func (b *limitedBuffer) Write(p []byte) (int, error) {
	if b.limit <= 0 {
		b.truncated = true
		return len(p), nil
	}
	remaining := b.limit - len(b.data)
	if remaining <= 0 {
		b.truncated = true
		return len(p), nil
	}
	if len(p) > remaining {
		b.data = append(b.data, p[:remaining]...)
		b.truncated = true
		return len(p), nil
	}
	b.data = append(b.data, p...)
	return len(p), nil
}

func (b *limitedBuffer) Bytes() []byte {
	return append([]byte{}, b.data...)
}

func (b *limitedBuffer) Truncated() bool {
	return b.truncated
}

var _ io.Writer = (*limitedBuffer)(nil)