wrenn-releases/internal/api/handlers_exec.go

package api

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"time"
	"unicode/utf8"

	"connectrpc.com/connect"

	"git.omukk.dev/wrenn/wrenn/pkg/auth"
	"git.omukk.dev/wrenn/wrenn/pkg/db"
	"git.omukk.dev/wrenn/wrenn/pkg/id"
	"git.omukk.dev/wrenn/wrenn/pkg/lifecycle"
	pb "git.omukk.dev/wrenn/wrenn/proto/hostagent/gen"
)

type execHandler struct {
	db   *db.Queries
	pool *lifecycle.HostClientPool
}

func newExecHandler(db *db.Queries, pool *lifecycle.HostClientPool) *execHandler {
	return &execHandler{db: db, pool: pool}
}

type execRequest struct {
	Cmd        string            `json:"cmd"`
	Args       []string          `json:"args"`
	TimeoutSec int32             `json:"timeout_sec"`
	Background bool              `json:"background"`
	Tag        string            `json:"tag"`
	Envs       map[string]string `json:"envs"`
	Cwd        string            `json:"cwd"`
}

type execResponse struct {
	SandboxID  string `json:"sandbox_id"`
	Cmd        string `json:"cmd"`
	Stdout     string `json:"stdout"`
	Stderr     string `json:"stderr"`
	ExitCode   int32  `json:"exit_code"`
	DurationMs int64  `json:"duration_ms"`
	// Encoding is "utf-8" for text output, "base64" for binary output.
	Encoding string `json:"encoding"`
}

type backgroundExecResponse struct {
	SandboxID string `json:"sandbox_id"`
	Cmd       string `json:"cmd"`
	PID       uint32 `json:"pid"`
	Tag       string `json:"tag"`
}

// Exec handles POST /v1/capsules/{id}/exec.
func (h *execHandler) Exec(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	ac := auth.MustFromContext(ctx)

	sb, sandboxID, sandboxIDStr, ok := requireRunningSandbox(w, r, h.db, ac.TeamID)
	if !ok {
		return
	}

	var req execRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid JSON body")
		return
	}

	if req.Cmd == "" {
		writeError(w, http.StatusBadRequest, "invalid_request", "cmd is required")
		return
	}

	agent, err := agentForHost(ctx, h.db, h.pool, sb.HostID)
	if err != nil {
		writeError(w, http.StatusServiceUnavailable, "host_unavailable", "sandbox host is not reachable")
		return
	}

	// Background mode: start process and return immediately.
	if req.Background {
		tag := req.Tag
		if tag == "" {
			tag = "proc-" + id.NewPtyTag()
		}

		bgResp, err := agent.StartBackground(ctx, connect.NewRequest(&pb.StartBackgroundRequest{
			SandboxId: sandboxIDStr,
			Tag:       tag,
			Cmd:       req.Cmd,
			Args:      req.Args,
			Envs:      req.Envs,
			Cwd:       req.Cwd,
		}))
		if err != nil {
			status, code, msg := agentErrToHTTP(err)
			writeError(w, status, code, msg)
			return
		}

		updateLastActive(h.db, sandboxID, sandboxIDStr)

		writeJSON(w, http.StatusAccepted, backgroundExecResponse{
			SandboxID: sandboxIDStr,
			Cmd:       req.Cmd,
			PID:       bgResp.Msg.Pid,
			Tag:       bgResp.Msg.Tag,
		})
		return
	}

	start := time.Now()

	resp, err := agent.Exec(ctx, connect.NewRequest(&pb.ExecRequest{
		SandboxId:  sandboxIDStr,
		Cmd:        req.Cmd,
		Args:       req.Args,
		TimeoutSec: req.TimeoutSec,
	}))
	if err != nil {
		status, code, msg := agentErrToHTTP(err)
		writeError(w, status, code, msg)
		return
	}

	duration := time.Since(start)

	updateLastActive(h.db, sandboxID, sandboxIDStr)

	stdout := resp.Msg.Stdout
	stderr := resp.Msg.Stderr

	encoding := "utf-8"
	stdoutStr, stderrStr := string(stdout), string(stderr)
	if !utf8.Valid(stdout) || !utf8.Valid(stderr) {
		encoding = "base64"
		stdoutStr = base64.StdEncoding.EncodeToString(stdout)
		stderrStr = base64.StdEncoding.EncodeToString(stderr)
	}

	writeJSON(w, http.StatusOK, execResponse{
		SandboxID:  sandboxIDStr,
		Cmd:        req.Cmd,
		Stdout:     stdoutStr,
		Stderr:     stderrStr,
		ExitCode:   resp.Msg.ExitCode,
		DurationMs: duration.Milliseconds(),
		Encoding:   encoding,
	})
}