135 lines
3.6 KiB
Go
135 lines
3.6 KiB
Go
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 {
|
|
buffer bytes.Buffer
|
|
exceeded bool
|
|
}
|
|
|
|
func (b *outputBuffer) Write(p []byte) (int, error) {
|
|
if len(p) > maxOutput-b.buffer.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.buffer.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))
|
|
}
|
|
}
|