foundry-core/internal/codex/canvas_worker.go

package codex

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

func executeCanvasTask(w http.ResponseWriter, r *http.Request, input CanvasTask, model CanvasModel) {
	imageOutput := input.OutputModality == "image"
	if input.OutputModality != "" && input.OutputModality != "text" && !imageOutput {
		fail(w, 400, "invalid_ai_request")
		return
	}
	if imageOutput != (len(model.OutputModalities) == 1 && model.OutputModalities[0] == "image") {
		fail(w, 422, "model_incompatible")
		return
	}
	duration := 90 * time.Second
	if imageOutput {
		duration = 180 * time.Second
	}
	ctx, cancel := context.WithTimeout(r.Context(), duration)
	defer cancel()
	job, err := os.MkdirTemp("", "foundry-canvas-")
	if err != nil {
		fail(w, 503, "codex_unavailable")
		return
	}
	defer os.RemoveAll(job)
	home := filepath.Join(job, "home")
	if os.Mkdir(home, 0700) != nil || os.Symlink("/home/node/.codex/auth.json", filepath.Join(home, "auth.json")) != nil {
		fail(w, 503, "codex_unavailable")
		return
	}
	args := []string{"--ask-for-approval", "never", "--disable", "shell_tool", "--disable", "apps", "--disable", "plugins", "--disable", "multi_agent", "--disable", "hooks", "--disable", "memories", "-c", `web_search="disabled"`}
	if !imageOutput {
		args = append(args, "--disable", "image_generation")
	}
	args = append(args, "exec", "--json", "--sandbox", "read-only", "--skip-git-repo-check", "--ephemeral", "--ignore-user-config", "--ignore-rules", "--model", model.Backend)
	total := 0
	for i, image := range input.Images {
		total += len(image.Data)
		if total > CanvasImageBytesLimit || len(image.Data) == 0 || http.DetectContentType(image.Data) != image.MIME {
			fail(w, 400, "invalid_ai_request")
			return
		}
		extension := map[string]string{"image/png": "png", "image/jpeg": "jpeg", "image/webp": "webp", "image/gif": "gif"}[image.MIME]
		if extension == "" {
			fail(w, 400, "invalid_ai_request")
			return
		}
		path := filepath.Join(job, fmt.Sprintf("input-%d.%s", i, extension))
		if os.WriteFile(path, image.Data, 0600) != nil {
			fail(w, 503, "codex_unavailable")
			return
		}
		args = append(args, "--image", path)
	}
	args = append(args, "-")
	prompt := input.Prompt
	if imageOutput {
		prompt = "Use the built-in image_generation tool to generate exactly one raster image following the supplied canvas instruction. Do not use shell, SVG, external URLs or substitute a textual description. Return a brief confirmation after generating the image.\n" + prompt
	}
	if len(prompt) > CanvasInputLimit {
		fail(w, 413, "context_too_large")
		return
	}
	cmd := exec.CommandContext(ctx, executable, args...)
	cmd.Dir = "/work"
	cmd.Env = []string{"PATH=/usr/local/bin:/usr/bin:/bin", "HOME=/home/node", "CODEX_HOME=" + home, "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
	output := outputBuffer{limit: 2 << 20}
	cmd.Stdout = &output
	if err = cmd.Run(); err != nil {
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			fail(w, 504, "codex_timeout")
		} else {
			fail(w, 502, "generation_failed")
		}
		return
	}
	if output.exceeded {
		fail(w, 502, "generation_failed")
		return
	}
	text, completed := canvasAgentText(output.buffer.Bytes())
	if !completed {
		fail(w, 502, "generation_failed")
		return
	}
	result := CanvasOutput{Kind: "text", Text: text}
	if imageOutput {
		image, err := readCanvasGeneratedImage(filepath.Join(home, "generated_images"))
		if err != nil {
			fail(w, 502, "generation_failed")
			return
		}
		result = CanvasOutput{Kind: "image", Image: image}
	} else if len(text) > maxOutput {
		fail(w, 502, "generation_failed")
		return
	}
	writeJSON(w, 200, map[string]any{"output": result})
}
func canvasAgentText(raw []byte) (string, bool) {
	result := ""
	completed := false
	for _, line := range bytes.Split(bytes.TrimSpace(raw), []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 {
			return "", false
		}
		if event.Type == "turn.failed" || event.Type == "error" {
			return "", false
		}
		if event.Type == "item.completed" && event.Item.Type == "agent_message" {
			result = event.Item.Text
		}
		if event.Type == "turn.completed" {
			completed = true
		}
	}
	return result, completed
}
func readCanvasGeneratedImage(root string) (*CanvasImage, error) {
	var result *CanvasImage
	err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if entry.Type()&os.ModeSymlink != 0 {
			return errors.New("invalid_artifact")
		}
		if entry.IsDir() {
			return nil
		}
		info, err := entry.Info()
		if err != nil {
			return err
		}
		if !info.Mode().IsRegular() {
			return errors.New("invalid_artifact")
		}
		ext := strings.ToLower(filepath.Ext(path))
		if ext != ".png" && ext != ".jpeg" && ext != ".jpg" && ext != ".webp" && ext != ".gif" {
			return nil
		}
		if result != nil || info.Size() <= 0 || info.Size() > CanvasImageBytesLimit {
			return errors.New("invalid_artifact")
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		mime := http.DetectContentType(data)
		switch mime {
		case "image/png", "image/jpeg", "image/webp", "image/gif":
		default:
			return errors.New("invalid_artifact")
		}
		result = &CanvasImage{MIME: mime, Data: data}
		return nil
	})
	if err != nil {
		return nil, err
	}
	if result == nil {
		return nil, errors.New("image_missing")
	}
	return result, nil
}