weatherreporter/internal/adapters/scriptorium/runner.go

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

import (
	"context"
	"fmt"
	"io"
	"os/exec"
	"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")
	}
	execution, err := r.execute(ctx, r.renderArgs(req))
	if err != nil {
		return nil, fmt.Errorf("run scriptorium render: %w", err)
	}
	result := &RenderResult{
		Command:         execution.argv(),
		Stdout:          string(execution.result.Stdout),
		Stderr:          string(execution.result.Stderr),
		StdoutTruncated: execution.result.StdoutTruncated,
		StderrTruncated: execution.result.StderrTruncated,
		ExitCode:        execution.result.ExitCode,
	}
	if execution.result.ExitCode != 0 {
		return result, fmt.Errorf("scriptorium render exited with code %d: %s", execution.result.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")
	}
	execution, err := r.execute(ctx, r.runArgs(req))
	if err != nil {
		return nil, fmt.Errorf("run scriptorium: %w", err)
	}
	result := &RunResult{
		Command:         execution.argv(),
		Stdout:          string(execution.result.Stdout),
		Stderr:          string(execution.result.Stderr),
		StdoutTruncated: execution.result.StdoutTruncated,
		StderrTruncated: execution.result.StderrTruncated,
		ExitCode:        execution.result.ExitCode,
		OutputPath:      req.OutputPath,
	}
	if execution.result.ExitCode != 0 {
		return result, fmt.Errorf("scriptorium run exited with code %d: %s", execution.result.ExitCode, result.Stderr)
	}
	return result, nil
}

type execution struct {
	binary string
	args   []string
	result CommandResult
}

func (r Runner) execute(ctx context.Context, args []string) (execution, error) {
	binary := r.Binary
	if binary == "" {
		binary = "scriptorium"
	}
	commands := r.Commands
	if commands == nil {
		commands = ExecRunner{}
	}
	result, err := commands.Run(ctx, binary, args, r.Timeout)
	if err != nil {
		return execution{}, err
	}
	return execution{binary: binary, args: args, result: result}, nil
}

func (e execution) argv() []string {
	return append([]string{e.binary}, e.args...)
}

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
}

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)