byre/internal/agents/imagegen.go

package agents

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"g.fc.run/ai/byre/internal/codex"
	"g.fc.run/ai/byre/internal/mcp"
)

const imageGenGuidance = "You make images for another agent. Call your image_gen tool exactly once with the prompt you are given, unchanged, then reply with only the word done."

func (m *Manager) generateImageTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "generate_image",
		Description: "Generate an image (a photo, illustration, texture, sprite, mockup) with OpenAI's image model through byre's Codex sign-in, saved as a PNG in the task's media folder. " +
			"The reply gives its path: show it inline by embedding it in your final message as a Markdown image, ![what it shows](<path>), or copy it into the worktree when the project uses it. " +
			"Use it whenever the human asks for a picture, rather than drawing one in SVG or HTML. It takes a minute or two.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"prompt": {"type": "string", "description": "What to draw: subject, style, composition, lighting, and any exact text."}
			},
			"required": ["prompt"]
		}`),
		Handler: mcp.Typed(func(ctx context.Context, args struct {
			Prompt string `json:"prompt"`
		}) (string, error) {
			return m.generateImage(ctx, l, args.Prompt)
		}),
	}
}

func (m *Manager) generateImage(ctx context.Context, l *live, prompt string) (string, error) {
	if strings.TrimSpace(prompt) == "" {
		return "", fmt.Errorf("generate_image needs a prompt")
	}

	media := MediaDir(l.dir)
	err := os.MkdirAll(m.codexHome, 0o700)
	if err != nil {
		return "", err
	}
	s, err := codex.Start(ctx, codex.Options{
		Cwd:      media,
		Env:      []string{"CODEX_HOME=" + m.codexHome},
		Guidance: imageGenGuidance,
	})
	if err != nil {
		return "", fmt.Errorf("start codex: %w", err)
	}
	defer s.Close()

	_, err = s.Send(ctx, prompt)
	if err != nil {
		return "", fmt.Errorf("start codex turn: %w", err)
	}

	saved := ""
	for {
		select {
		case <-ctx.Done():
			s.Kill()
			return "", fmt.Errorf("no image after %s", mcp.Timeout(ctx))
		case ev, ok := <-s.Events():
			if !ok {
				return "", fmt.Errorf("codex exited before the image was made")
			}
			switch ev.Method {
			case codex.MethodItemCompleted:
				n := codex.ItemCompletedNotification{}
				err := json.Unmarshal(ev.Params, &n)
				if err != nil {
					return "", err
				}
				if n.Item.Type != codex.ThreadItemTypeImageGeneration {
					continue
				}
				if n.Item.Failure != nil {
					return "", fmt.Errorf("image generation failed: %s (limit %s)", n.Item.Failure.Type, n.Item.Failure.LimitID)
				}
				saved = string(n.Item.SavedPath)
			case codex.MethodTurnCompleted:
				n := codex.TurnCompletedNotification{}
				err := json.Unmarshal(ev.Params, &n)
				if err != nil {
					return "", err
				}
				if n.Turn.Error != nil {
					return "", fmt.Errorf("codex: %s", n.Turn.Error.Message)
				}
				if saved == "" {
					return "", fmt.Errorf("codex finished without making an image")
				}
				path, err := copyToMedia(media, saved)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Saved %s", path), nil
			}
		}
	}
}

func copyToMedia(media, src string) (string, error) {
	data, err := os.ReadFile(src)
	if err != nil {
		return "", err
	}
	err = os.MkdirAll(media, 0o755)
	if err != nil {
		return "", err
	}
	dest := filepath.Join(media, filepath.Base(src))
	return dest, os.WriteFile(dest, data, 0o644)
}