Add project-scoped Codex execution over a private socket

internal/codex/worker.go

@@ -0,0 +1,134 @@
+package codex
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"os"
+	"os/exec"
+	"strings"
+	"syscall"
+	"time"
+)
+
+const executable = "/opt/codex/node_modules/.bin/codex"
+
+var errOutputLimit = errors.New("Codex output limit exceeded")
+
+type outputBuffer struct {
+	bytes.Buffer
+	exceeded bool
+}
+
+func (b *outputBuffer) Write(p []byte) (int, error) {
+	if len(p) > maxOutput-b.Len() {
+		b.exceeded = true
+		return 0, errOutputLimit
+	}
+	return b.Buffer.Write(p)
+}
+
+func ServeWorker(socket string) error {
+	syscall.Umask(0077)
+	if err := os.Remove(socket); err != nil && !errors.Is(err, os.ErrNotExist) {
+		return err
+	}
+	listener, err := net.Listen("unix", socket)
+	if err != nil {
+		return err
+	}
+	defer listener.Close()
+	if err := os.Chmod(socket, 0600); err != nil {
+		return err
+	}
+	slot := make(chan struct{}, 1)
+	mux := http.NewServeMux()
+	mux.HandleFunc("POST /exec", func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case slot <- struct{}{}:
+			defer func() { <-slot }()
+		default:
+			w.Header().Set("Retry-After", "5")
+			fail(w, 429, "codex_busy")
+			return
+		}
+		input, ok := readTask(w, r, 6*maxRequest)
+		if ok {
+			execute(w, r, input.Prompt)
+		}
+	})
+	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, ReadTimeout: 10 * time.Second, WriteTimeout: timeout + 15*time.Second}
+	return server.Serve(listener)
+}
+
+func execute(w http.ResponseWriter, r *http.Request, prompt string) {
+	ctx, cancel := context.WithTimeout(r.Context(), timeout)
+	defer cancel()
+	cmd := exec.CommandContext(ctx, executable, "--ask-for-approval", "never",
+		"--disable", "shell_tool", "--disable", "apps", "--disable", "plugins",
+		"--disable", "multi_agent", "--disable", "hooks", "--disable", "memories",
+		"-c", `web_search="disabled"`, "exec", "--json", "--sandbox", "read-only", "--skip-git-repo-check", "--ephemeral", "-")
+	cmd.Dir = "/work"
+	cmd.Env = []string{"PATH=/usr/local/bin:/usr/bin:/bin", "HOME=/home/node", "CODEX_HOME=/home/node/.codex", "SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt"}
+	cmd.Stdin = strings.NewReader(prompt)
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	cmd.Cancel = func() error {
+		err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+		if errors.Is(err, syscall.ESRCH) {
+			return os.ErrProcessDone
+		}
+		return err
+	}
+	cmd.WaitDelay = 2 * time.Second
+	var output outputBuffer
+	cmd.Stdout, cmd.Stderr = &output, io.Discard
+	err := cmd.Run()
+	switch {
+	case output.exceeded:
+		fail(w, 502, "codex_output_too_large")
+	case errors.Is(ctx.Err(), context.DeadlineExceeded):
+		fail(w, 504, "codex_timeout")
+	case err != nil:
+		fail(w, 502, "codex_failed")
+	default:
+		var result string
+		var completed bool
+		for _, line := range bytes.Split(bytes.TrimSpace(output.Bytes()), []byte("\n")) {
+			var event struct {
+				Type string `json:"type"`
+				Item struct {
+					Type string `json:"type"`
+					Text string `json:"text"`
+				} `json:"item"`
+			}
+			if json.Unmarshal(line, &event) != nil {
+				fail(w, 502, "invalid_codex_response")
+				return
+			}
+			if event.Type == "item.completed" && event.Item.Type == "agent_message" {
+				result = event.Item.Text
+			}
+			if event.Type == "turn.failed" || event.Type == "error" {
+				fail(w, 502, "codex_failed")
+				return
+			}
+			if event.Type == "turn.completed" {
+				completed = true
+			}
+		}
+		if !completed || result == "" {
+			fail(w, 502, "codex_failed")
+			return
+		}
+		response, _ := json.Marshal(map[string]string{"result": result})
+		if len(response) >= maxOutput {
+			fail(w, 502, "codex_output_too_large")
+			return
+		}
+		writeJSON(w, 200, json.RawMessage(response))
+	}
+}