iop/apps/node/internal/adapters/cli/oneshot.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

func (c *CLI) executeOneShot(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf, sink runtime.EventSink) error {
	prompt := extractPrompt(spec.Input)
	args := append([]string{}, profile.Args...)
	if prompt != "" {
		args = append(args, prompt)
	}
	_, err := c.executeCommand(ctx, spec, profile, args, prompt, sink)
	return err
}

func (c *CLI) executeCommand(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf, args []string, prompt string, sink runtime.EventSink) (string, error) {
	cmd, err := buildCmd(ctx, profile.Command, args, profile.Env, spec.Workspace)
	if err != nil {
		return "", emitReturnedError(ctx, sink, spec.RunID, err)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", emitReturnedError(ctx, sink, spec.RunID, fmt.Errorf("cli adapter: stdout pipe: %w", err))
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return "", emitReturnedError(ctx, sink, spec.RunID, fmt.Errorf("cli adapter: stderr pipe: %w", err))
	}

	if err := cmd.Start(); err != nil {
		// The context can be cancelled or time out before the process starts,
		// which surfaces as a start error. Honor the cancellation contract with a
		// single cancelled event and ErrRunCancelled instead of a start error.
		if ctx.Err() != nil {
			_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeCancelled,
				Message:   cancelEventForContext(ctx.Err()),
				Timestamp: time.Now(),
			})
			return "", runtime.ErrRunCancelled
		}
		return "", emitReturnedError(ctx, sink, spec.RunID, fmt.Errorf("cli adapter: start %q: %w", profile.Command, err))
	}

	// Closing the stdout/stderr pipes on cancellation unblocks pending reads
	// promptly. CommandContext only signals the direct child, so a forked
	// grandchild that keeps the pipes open would otherwise stall the read until
	// it exits on its own; closing the reader ends the read-vs-wait race.
	readsDone := make(chan struct{})
	defer close(readsDone)
	go func() {
		select {
		case <-ctx.Done():
			_ = stdout.Close()
			_ = stderr.Close()
		case <-readsDone:
		}
	}()

	_ = sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     spec.RunID,
		Type:      runtime.EventTypeStart,
		Timestamp: time.Now(),
	})

	var errBuf strings.Builder
	var outBuf strings.Builder
	combinedOutput := func() string {
		return outBuf.String() + errBuf.String()
	}
	stderrDone := make(chan error, 1)
	go func() {
		_, err := io.Copy(&errBuf, stderr)
		stderrDone <- err
	}()

	var (
		outputTokens int
		readErr      error
	)
	if reg, ok := jsonEmitters[profile.OutputFormat]; ok {
		outputTokens, readErr = driveJSONLines(ctx, stdout, sink, spec.RunID, &outBuf, reg.emitter, reg.scanBufMax)
	} else {
		outputTokens, readErr = emitStdoutChunks(ctx, stdout, sink, spec.RunID, &outBuf)
	}
	if readErr != nil {
		_ = <-stderrDone
		_ = cmd.Wait()
		// A context cancellation can close the stdout pipe before cmd.Wait()
		// observes the process exit, surfacing as a read error. Treat that as a
		// cancellation so the user-cancel/timeout contract wins over the generic
		// read-stdout error, emitting exactly one cancelled event.
		if ctx.Err() != nil {
			_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeCancelled,
				Message:   cancelEventForContext(ctx.Err()),
				Timestamp: time.Now(),
			})
			return combinedOutput(), runtime.ErrRunCancelled
		}
		_ = sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     spec.RunID,
			Type:      runtime.EventTypeError,
			Error:     fmt.Sprintf("read stdout: %v", readErr),
			Timestamp: time.Now(),
		})
		return combinedOutput(), fmt.Errorf("cli adapter: read stdout: %w", readErr)
	}

	waitErr := cmd.Wait()
	stderrErr := <-stderrDone
	if waitErr != nil {
		if ctx.Err() != nil {
			_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeCancelled,
				Message:   cancelEventForContext(ctx.Err()),
				Timestamp: time.Now(),
			})
			return combinedOutput(), runtime.ErrRunCancelled
		}
		_ = sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     spec.RunID,
			Type:      runtime.EventTypeError,
			Error:     fmt.Sprintf("command failed: %v — %s", waitErr, errBuf.String()),
			Timestamp: time.Now(),
		})
		return combinedOutput(), fmt.Errorf("cli adapter: command exited with error: %w", waitErr)
	}
	if stderrErr != nil && !errors.Is(stderrErr, os.ErrClosed) {
		_ = sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     spec.RunID,
			Type:      runtime.EventTypeError,
			Error:     fmt.Sprintf("read stderr: %v", stderrErr),
			Timestamp: time.Now(),
		})
		return combinedOutput(), fmt.Errorf("cli adapter: read stderr: %w", stderrErr)
	}

	return combinedOutput(), sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:   spec.RunID,
		Type:    runtime.EventTypeComplete,
		Message: "cli execution complete",
		Usage: &runtime.UsageStats{
			InputTokens:  len(strings.Fields(prompt)),
			OutputTokens: outputTokens,
		},
		Timestamp: time.Now(),
	})
}

func emitStdoutChunks(ctx context.Context, stdout io.Reader, sink runtime.EventSink, runID string, outBuf *strings.Builder) (int, error) {
	buf := make([]byte, 4096)
	outputTokens := 0

	for {
		n, err := stdout.Read(buf)
		if n > 0 {
			delta := string(buf[:n])
			outBuf.WriteString(delta)
			outputTokens += len(strings.Fields(delta))
			_ = sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:     runID,
				Type:      runtime.EventTypeDelta,
				Delta:     delta,
				Timestamp: time.Now(),
			})
		}
		if errors.Is(err, io.EOF) {
			return outputTokens, nil
		}
		if err != nil {
			return outputTokens, err
		}
	}
}