kandev/apps/backend/internal/worktree/script_message_handler.go

package worktree

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sort"
	"sync"
	"time"

	"go.uber.org/zap"

	"github.com/kandev/kandev/internal/common/logger"
	"github.com/kandev/kandev/internal/common/shellexec"
	"github.com/kandev/kandev/internal/task/models"
)

// ScriptExecutionRequest contains parameters for executing a setup or cleanup script.
type ScriptExecutionRequest struct {
	SessionID    string
	TaskID       string
	RepositoryID string
	Script       string
	WorkingDir   string
	ScriptType   string // "setup" or "cleanup"
	Env          map[string]string
}

// DefaultScriptMessageHandler manages script execution and streams output to session messages.
type DefaultScriptMessageHandler struct {
	logger      *logger.Logger
	taskService TaskService
	timeout     time.Duration
}

// TaskService interface for creating and updating messages.
type TaskService interface {
	CreateMessage(ctx context.Context, req *CreateMessageRequest) (*models.Message, error)
	UpdateMessage(ctx context.Context, message *models.Message) error
}

// CreateMessageRequest contains parameters for creating a message.
type CreateMessageRequest struct {
	TaskSessionID string
	TaskID        string
	TurnID        string
	Content       string
	AuthorType    string
	AuthorID      string
	RequestsInput bool
	Type          string
	Metadata      map[string]interface{}
}

// NewDefaultScriptMessageHandler creates a new DefaultScriptMessageHandler.
func NewDefaultScriptMessageHandler(
	log *logger.Logger,
	taskSvc TaskService,
	timeout time.Duration,
) *DefaultScriptMessageHandler {
	return &DefaultScriptMessageHandler{
		logger:      log.WithFields(zap.String("component", "script-message-handler")),
		taskService: taskSvc,
		timeout:     timeout,
	}
}

// ExecuteSetupScript executes a setup script and streams output to a session message.
// Returns an error if the script fails (non-zero exit code or timeout).
func (h *DefaultScriptMessageHandler) ExecuteSetupScript(ctx context.Context, req ScriptExecutionRequest) error {
	return h.executeScript(ctx, req, true)
}

// ExecuteCleanupScript executes a cleanup script and streams output to a session message.
// Returns nil even if the script fails (best-effort cleanup).
func (h *DefaultScriptMessageHandler) ExecuteCleanupScript(ctx context.Context, req ScriptExecutionRequest) error {
	err := h.executeScript(ctx, req, false)
	if err != nil {
		h.logger.Warn("cleanup script failed, continuing with deletion",
			zap.String("session_id", req.SessionID),
			zap.Error(err))
		return nil
	}
	return nil
}

// executeScript is the core implementation for script execution.
// Note: The parent context is intentionally not used - we create a detached context
// to prevent HTTP request timeouts from cancelling long-running scripts.
func (h *DefaultScriptMessageHandler) executeScript(_ context.Context, req ScriptExecutionRequest, failOnError bool) error {
	if h.taskService == nil {
		h.logger.Debug("script handler not fully configured, skipping",
			zap.String("script_type", req.ScriptType))
		return nil
	}

	// Create a detached context for script execution with its own timeout.
	// This prevents the HTTP request context from cancelling the script.
	scriptCtx, cancel := context.WithTimeout(context.Background(), h.timeout)
	defer cancel()

	// Create initial message (best-effort - session may not exist during cleanup)
	msg, err := h.createScriptMessage(scriptCtx, req)
	if err != nil {
		// For cleanup scripts, if session doesn't exist, run the script anyway without a message
		if req.ScriptType == "cleanup" {
			h.logger.Warn("failed to create cleanup script message, running script without message tracking",
				zap.String("session_id", req.SessionID),
				zap.Error(err))
			// Run script directly without message tracking
			return h.runScriptWithoutMessage(scriptCtx, req, failOnError)
		}
		// For setup scripts, this is a hard error
		return fmt.Errorf("failed to create script message: %w", err)
	}

	h.logger.Info("created script execution message",
		zap.String("message_id", msg.ID),
		zap.String("session_id", req.SessionID),
		zap.String("script_type", req.ScriptType))

	// Update message status to running
	msg.Metadata["status"] = "running"
	if err := h.taskService.UpdateMessage(scriptCtx, msg); err != nil {
		h.logger.Warn("failed to update message with running status",
			zap.String("message_id", msg.ID),
			zap.Error(err))
	}

	// Execute the script and capture output
	exitCode, scriptErr := h.runScriptWithOutput(scriptCtx, req, msg)

	// Update final status
	applyScriptFinalStatus(msg, exitCode, scriptErr)

	msg.Metadata["completed_at"] = time.Now().UTC().Format(time.RFC3339Nano)
	if updateErr := h.taskService.UpdateMessage(scriptCtx, msg); updateErr != nil {
		h.logger.Warn("failed to update message with final status",
			zap.String("message_id", msg.ID),
			zap.Error(updateErr))
	}

	h.logger.Info("script execution completed",
		zap.String("message_id", msg.ID),
		zap.Int("exit_code", exitCode),
		zap.Bool("success", exitCode == 0))

	return scriptExecutionError(failOnError, exitCode, scriptErr)
}

// applyScriptFinalStatus updates the message metadata and content based on script outcome.
func applyScriptFinalStatus(msg *models.Message, exitCode int, scriptErr error) {
	switch {
	case scriptErr != nil:
		msg.Metadata["status"] = "failed"
		msg.Metadata["error"] = scriptErr.Error()
		if msg.Content == "" {
			msg.Content = fmt.Sprintf("Script execution failed: %s", scriptErr.Error())
		} else {
			msg.Content += fmt.Sprintf("\n\nScript execution failed: %s", scriptErr.Error())
		}
	case exitCode == 0:
		msg.Metadata["status"] = "exited"
		msg.Metadata["exit_code"] = exitCode
		if msg.Content == "" {
			msg.Content = "Script completed successfully"
		}
	default:
		msg.Metadata["status"] = "failed"
		msg.Metadata["exit_code"] = exitCode
		if msg.Content == "" {
			msg.Content = fmt.Sprintf("Script failed with exit code: %d", exitCode)
		} else {
			msg.Content += fmt.Sprintf("\n\nScript failed with exit code: %d", exitCode)
		}
	}
}

// scriptExecutionError returns an error if failOnError is set and the script did not succeed.
func scriptExecutionError(failOnError bool, exitCode int, scriptErr error) error {
	if !failOnError {
		return nil
	}
	if scriptErr != nil {
		return scriptErr
	}
	if exitCode != 0 {
		return fmt.Errorf("script exited with code %d", exitCode)
	}
	return nil
}

// createScriptMessage creates the initial script execution message.
func (h *DefaultScriptMessageHandler) createScriptMessage(ctx context.Context, req ScriptExecutionRequest) (*models.Message, error) {
	metadata := map[string]interface{}{
		"script_type": req.ScriptType,
		"command":     req.Script,
		"status":      "starting",
		"started_at":  time.Now().UTC().Format(time.RFC3339Nano),
	}

	createReq := &CreateMessageRequest{
		TaskSessionID: req.SessionID,
		TaskID:        req.TaskID,
		Content:       "", // Will be populated with output
		AuthorType:    "agent",
		Type:          "script_execution",
		Metadata:      metadata,
	}

	return h.taskService.CreateMessage(ctx, createReq)
}

// runScriptWithOutput runs the script and captures output, streaming it to the message.
// The passed context should already have an appropriate timeout set.
func (h *DefaultScriptMessageHandler) runScriptWithOutput(ctx context.Context, req ScriptExecutionRequest, msg *models.Message) (int, error) {
	cmd := shellexec.CommandContext(ctx, shellexec.PosixSh, req.Script)
	cmd.Dir = req.WorkingDir
	cmd.Env = scriptProcessEnvironment(req.Env)

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return -1, fmt.Errorf("failed to create stdout pipe: %w", err)
	}
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return -1, fmt.Errorf("failed to create stderr pipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return -1, fmt.Errorf("failed to start script: %w", err)
	}

	h.logger.Info("script process started",
		zap.String("message_id", msg.ID),
		zap.String("command", req.Script))

	var outputBuf bytes.Buffer
	var wg sync.WaitGroup
	var mu sync.Mutex

	wg.Add(2)
	go h.streamPipeOutput(stdoutPipe, &outputBuf, &mu, msg, &wg)
	go h.streamPipeOutput(stderrPipe, &outputBuf, &mu, msg, &wg)
	wg.Wait()

	err = cmd.Wait()
	exitCode := 0
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
		} else {
			return -1, err
		}
	}
	return exitCode, nil
}

// streamPipeOutput reads from pipe and streams output into msg.Content incrementally.
func (h *DefaultScriptMessageHandler) streamPipeOutput(pipe io.Reader, outputBuf *bytes.Buffer, mu *sync.Mutex, msg *models.Message, wg *sync.WaitGroup) {
	defer wg.Done()
	buf := make([]byte, 1024)
	for {
		n, err := pipe.Read(buf)
		if n > 0 {
			mu.Lock()
			outputBuf.Write(buf[:n])
			msg.Content = outputBuf.String()
			mu.Unlock()
			if updateErr := h.taskService.UpdateMessage(context.Background(), msg); updateErr != nil {
				h.logger.Debug("failed to update message with output",
					zap.String("message_id", msg.ID),
					zap.Error(updateErr))
			}
		}
		if err != nil {
			if err != io.EOF {
				h.logger.Debug("error reading pipe output", zap.Error(err))
			}
			break
		}
	}
}

// runScriptWithoutMessage runs a script without message tracking (used when session is deleted).
// The passed context should already have an appropriate timeout set.
func (h *DefaultScriptMessageHandler) runScriptWithoutMessage(ctx context.Context, req ScriptExecutionRequest, failOnError bool) error {
	h.logger.Info("executing script without message tracking",
		zap.String("script_type", req.ScriptType),
		zap.String("command", req.Script))

	// Run script under the host shell (sh -c on Unix, bash/cmd on Windows).
	cmd := shellexec.CommandContext(ctx, shellexec.PosixSh, req.Script)
	cmd.Dir = req.WorkingDir
	cmd.Env = scriptProcessEnvironment(req.Env)

	// Capture output for logging
	output, err := cmd.CombinedOutput()

	// Get exit code
	exitCode := 0
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
		} else {
			h.logger.Error("script execution failed",
				zap.String("script_type", req.ScriptType),
				zap.Error(err),
				zap.String("output", string(output)))
			if failOnError {
				return err
			}
			return nil
		}
	}

	// Log the result
	if exitCode == 0 {
		h.logger.Info("script completed successfully",
			zap.String("script_type", req.ScriptType),
			zap.Int("exit_code", exitCode),
			zap.String("output", string(output)))
	} else {
		h.logger.Warn("script failed",
			zap.String("script_type", req.ScriptType),
			zap.Int("exit_code", exitCode),
			zap.String("output", string(output)))
		if failOnError {
			return fmt.Errorf("script exited with code %d", exitCode)
		}
	}

	return nil
}

func scriptProcessEnvironment(overrides map[string]string) []string {
	if len(overrides) == 0 {
		return nil
	}
	env := os.Environ()
	keys := make([]string, 0, len(overrides))
	for key := range overrides {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, key := range keys {
		env = append(env, key+"="+overrides[key])
	}
	return env
}