iop/apps/node/internal/workspace/command_executor.go

package workspace

import (
	"bytes"
	"context"
	"errors"
	"io"
	"math"
	"os"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"

	iop "iop/proto/gen/iop"
)

const (
	commandLaunchVersion      = 1
	commandLaunchRecordLimit  = 64 << 10
	commandLaunchPayloadLimit = 4 << 10
	commandShimArgument       = "__iop_workspace_command_shim"
	commandShimEnvironment    = "IOP_WORKSPACE_COMMAND_SHIM"
)

var (
	errCommandPlatformUnsupported = errors.New("workspace command execution is unsupported on this platform")
	errCommandLaunchInvalid       = errors.New("workspace command launch is invalid")
)

type commandTemplate struct {
	executable string
	args       []string
}

type commandKey struct {
	requestID  string
	toolCallID string
}

type commandExecution struct {
	mu              sync.Mutex
	cancelRequested bool
	finished        bool
	cancel          chan struct{}
	done            chan struct{}
	doneOnce        sync.Once
}

func newCommandExecution() *commandExecution {
	return &commandExecution{cancel: make(chan struct{}), done: make(chan struct{})}
}

func (e *commandExecution) requestCancel() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.finished {
		return false
	}
	if !e.cancelRequested {
		e.cancelRequested = true
		close(e.cancel)
	}
	return true
}

func (e *commandExecution) finish() {
	e.mu.Lock()
	e.finished = true
	e.mu.Unlock()
	e.doneOnce.Do(func() { close(e.done) })
}

// CommandInput contains only request identity and caller-selectable fields
// already closed by the workspace wire. Executable and argv never enter it.
type CommandInput struct {
	RequestID   string
	ToolCallID  string
	CommandID   string
	Environment map[string]string
	TimeoutMS   int64
}

// CancelResult is the stable result of addressing one active command by its
// exact request and tool-call identity.
type CancelResult struct {
	Status iop.WorkspaceStatus
	Code   iop.WorkspaceErrorCode
}

type commandLaunchRecord struct {
	Version     int      `json:"version"`
	Executable  string   `json:"executable"`
	Args        []string `json:"args,omitempty"`
	Environment []string `json:"environment,omitempty"`
	Device      uint64   `json:"device"`
	Inode       uint64   `json:"inode"`
}

type commandLaunchStatus struct {
	started bool
}

type commandProcess struct {
	wait     <-chan error
	launch   <-chan commandLaunchStatus
	pid      int
	exitCode func() int32
}

// ExecuteCommand resolves an admitted command id to one immutable operator
// template and owns its complete process/result lifecycle.
func (r *Runtime) ExecuteCommand(ctx context.Context, input CommandInput) (result Result) {
	observationStartedAt := time.Now()
	var correlation string
	defer func() {
		if result.DurationMS == 0 {
			result.DurationMS = time.Since(observationStartedAt).Milliseconds()
		}
		r.observeTool(correlation, iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND, result)
	}()
	if ctx == nil {
		ctx = context.Background()
	}
	r.lifetime.RLock()
	defer r.lifetime.RUnlock()

	req, template, environment, result := r.prepareCommand(input)
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
		return result
	}
	correlation = req.correlation
	if ctx.Err() != nil {
		return Result{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, ExitCode: -1}
	}
	key := commandKey{requestID: input.RequestID, toolCallID: input.ToolCallID}
	execution := newCommandExecution()
	if !r.registerCommand(req, key, execution) {
		return invalidCommandResult()
	}
	defer func() {
		execution.finish()
		r.commandsMu.Lock()
		if r.activeCommands[key] == execution {
			delete(r.activeCommands, key)
		}
		r.commandsMu.Unlock()
	}()

	output := newCommandOutput(req.maxOutput)
	record := commandLaunchRecord{
		Version: commandLaunchVersion, Executable: template.executable,
		Args: append([]string(nil), template.args...), Environment: environment,
		Device: req.entry.device, Inode: req.entry.inode,
	}

	execution.mu.Lock()
	if execution.cancelRequested {
		execution.mu.Unlock()
		return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, -1, 0, output)
	}
	startedAt := time.Now()
	process, err := startCommandProcess(record, req.entry.directory, output)
	execution.mu.Unlock()
	if err != nil {
		if errors.Is(err, errCommandPlatformUnsupported) {
			return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSUPPORTED, -1, time.Since(startedAt), output)
		}
		return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, -1, time.Since(startedAt), output)
	}
	return awaitCommand(ctx, execution, process, time.Duration(input.TimeoutMS)*time.Millisecond, startedAt, output)
}

func (r *Runtime) prepareCommand(input CommandInput) (*Request, commandTemplate, []string, Result) {
	if !validRequestID(input.RequestID) || !validRequestID(input.ToolCallID) || strings.TrimSpace(input.CommandID) == "" || input.CommandID != strings.TrimSpace(input.CommandID) {
		return nil, commandTemplate{}, nil, invalidCommandResult()
	}
	req, err := r.Request(input.RequestID)
	if err != nil {
		return nil, commandTemplate{}, nil, failureFor(err)
	}
	if _, allowed := req.operations[iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND]; !allowed {
		return nil, commandTemplate{}, nil, Result{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSUPPORTED}
	}
	if _, allowed := slices.BinarySearch(req.commandIDs, input.CommandID); !allowed {
		return nil, commandTemplate{}, nil, invalidCommandResult()
	}
	if input.TimeoutMS <= 0 || input.TimeoutMS > req.maxCommandTimeout || input.TimeoutMS > math.MaxInt64/int64(time.Millisecond) {
		return nil, commandTemplate{}, nil, invalidCommandResult()
	}
	template, ok := req.entry.commands[input.CommandID]
	if !ok {
		return nil, commandTemplate{}, nil, invalidCommandResult()
	}
	environment, ok := buildMinimalEnvironment(input.Environment, req.entry.environment)
	if !ok {
		return nil, commandTemplate{}, nil, invalidCommandResult()
	}
	return req, template, environment, Result{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
}

func (r *Runtime) registerCommand(req *Request, key commandKey, execution *commandExecution) bool {
	if req == nil {
		return false
	}
	req.mu.Lock()
	defer req.mu.Unlock()
	if req.cleaning {
		return false
	}
	r.commandsMu.Lock()
	defer r.commandsMu.Unlock()
	if _, duplicate := r.activeCommands[key]; duplicate {
		return false
	}
	if _, cancelled := r.cancelledCommands[key]; cancelled {
		return false
	}
	r.activeCommands[key] = execution
	return true
}

// Cancel requests process-group cancellation for one exact active command.
// Repeated requests remain idempotent for the open request lifecycle.
func (r *Runtime) Cancel(requestID, toolCallID string) CancelResult {
	if !validRequestID(requestID) || !validRequestID(toolCallID) {
		return CancelResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST}
	}
	key := commandKey{requestID: requestID, toolCallID: toolCallID}
	r.commandsMu.Lock()
	execution := r.activeCommands[key]
	_, alreadyCancelled := r.cancelledCommands[key]
	if alreadyCancelled {
		r.commandsMu.Unlock()
		return CancelResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED}
	}
	if execution == nil || !execution.requestCancel() {
		r.commandsMu.Unlock()
		return CancelResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_FOUND}
	}
	r.cancelledCommands[key] = struct{}{}
	r.commandsMu.Unlock()
	return CancelResult{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED}
}

func awaitCommand(ctx context.Context, execution *commandExecution, process *commandProcess, timeout time.Duration, startedAt time.Time, output *commandOutput) Result {
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	var (
		launchKnown    bool
		launched       bool
		waitDone       bool
		waitErr        error
		terminalStatus iop.WorkspaceStatus
		terminalCode   iop.WorkspaceErrorCode
		terminated     bool
	)
	for {
		if terminalStatus != iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED && !terminated {
			terminateProcessGroup(process.pid)
			terminated = true
		}
		if waitDone && terminalStatus != iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED {
			return terminalCommandResult(terminalStatus, terminalCode, -1, time.Since(startedAt), output)
		}
		if waitDone && launchKnown {
			if !launched {
				return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, -1, time.Since(startedAt), output)
			}
			exitCode := process.exitCode()
			if waitErr == nil && exitCode == 0 {
				return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSPECIFIED, 0, time.Since(startedAt), output)
			}
			return terminalCommandResult(iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL, exitCode, time.Since(startedAt), output)
		}

		select {
		case status := <-process.launch:
			launchKnown = true
			launched = status.started
			process.launch = nil
			if !launched && terminalStatus == iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED {
				terminalStatus = iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR
				terminalCode = iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL
			}
		case waitErr = <-process.wait:
			waitDone = true
			process.wait = nil
		case <-timer.C:
			if terminalStatus == iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED {
				terminalStatus = iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT
				terminalCode = iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT
			}
			timer.Stop()
		case <-execution.cancel:
			if terminalStatus == iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED {
				terminalStatus = iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED
				terminalCode = iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED
			}
			execution.cancel = nil
		case <-ctx.Done():
			if terminalStatus == iop.WorkspaceStatus_WORKSPACE_STATUS_UNSPECIFIED {
				terminalStatus = iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED
				terminalCode = iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED
			}
			ctx = context.Background()
		}
	}
}

func terminalCommandResult(status iop.WorkspaceStatus, code iop.WorkspaceErrorCode, exitCode int32, duration time.Duration, output *commandOutput) Result {
	stdout, stderr, truncated := output.snapshot()
	return Result{
		Status: status, Code: code, Stdout: stdout, Stderr: stderr,
		ExitCode: exitCode, Truncated: truncated, DurationMS: duration.Milliseconds(),
	}
}

func invalidCommandResult() Result {
	return Result{Status: iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR, Code: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST, ExitCode: -1}
}

func buildMinimalEnvironment(input map[string]string, allowlist map[string]struct{}) ([]string, bool) {
	if len(input) == 0 {
		return []string{}, true
	}
	names := make([]string, 0, len(input))
	total := 0
	for name, value := range input {
		if !validEnvironmentName(name) || name == commandShimEnvironment || strings.IndexByte(value, 0) >= 0 {
			return nil, false
		}
		if _, allowed := allowlist[name]; !allowed {
			return nil, false
		}
		total += len(name) + len(value) + 1
		if total > commandLaunchPayloadLimit {
			return nil, false
		}
		names = append(names, name)
	}
	sort.Strings(names)
	environment := make([]string, 0, len(names))
	for _, name := range names {
		environment = append(environment, name+"="+input[name])
	}
	return environment, true
}

func validEnvironmentName(name string) bool {
	if name == "" {
		return false
	}
	for index := 0; index < len(name); index++ {
		value := name[index]
		if index == 0 {
			if (value >= 'a' && value <= 'z') || (value >= 'A' && value <= 'Z') || value == '_' {
				continue
			}
			return false
		}
		if (value >= 'a' && value <= 'z') || (value >= 'A' && value <= 'Z') || (value >= '0' && value <= '9') || value == '_' {
			continue
		}
		return false
	}
	return true
}

type commandOutput struct {
	mu        sync.Mutex
	remaining int64
	truncated bool
	stdout    bytes.Buffer
	stderr    bytes.Buffer
}

type commandOutputWriter struct {
	output *commandOutput
	stderr bool
}

func newCommandOutput(limit int64) *commandOutput {
	return &commandOutput{remaining: limit}
}

func (o *commandOutput) writer(stderr bool) io.Writer {
	return commandOutputWriter{output: o, stderr: stderr}
}

func (w commandOutputWriter) Write(data []byte) (int, error) {
	w.output.mu.Lock()
	defer w.output.mu.Unlock()
	retained := int64(len(data))
	if retained > w.output.remaining {
		retained = w.output.remaining
		w.output.truncated = true
	}
	if retained < int64(len(data)) {
		w.output.truncated = true
	}
	if retained > 0 {
		if w.stderr {
			_, _ = w.output.stderr.Write(data[:retained])
		} else {
			_, _ = w.output.stdout.Write(data[:retained])
		}
		w.output.remaining -= retained
	}
	return len(data), nil
}

func (o *commandOutput) snapshot() ([]byte, []byte, bool) {
	o.mu.Lock()
	defer o.mu.Unlock()
	return bytes.Clone(o.stdout.Bytes()), bytes.Clone(o.stderr.Bytes()), o.truncated
}

// RunCommandShim must run before Cobra parsing. It recognizes only the exact
// internal invocation and otherwise leaves normal CLI behavior untouched.
func RunCommandShim(args []string) (bool, int) {
	if len(args) != 2 || args[1] != commandShimArgument || commandShimMarker() != "1" {
		return false, 0
	}
	return true, runCommandShim()
}

func commandShimMarker() string {
	return os.Getenv(commandShimEnvironment)
}