diple/ai_codex.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type CodexCLIProvider struct {
	command   string
	model     string
	timeout   time.Duration
	workspace string
	mu        sync.Mutex
	status    AIProviderStatus
	statusAt  time.Time
}

var errAIOutputLimit = errors.New("AI provider output limit exceeded")

type limitedBuffer struct {
	buffer bytes.Buffer
	limit  int
}

func (w *limitedBuffer) Write(value []byte) (int, error) {
	remaining := w.limit - w.buffer.Len()
	if remaining <= 0 {
		return 0, errAIOutputLimit
	}
	if len(value) > remaining {
		_, _ = w.buffer.Write(value[:remaining])
		return remaining, errAIOutputLimit
	}
	return w.buffer.Write(value)
}

func NewCodexCLIProvider(config AIConfig, workspace string) *CodexCLIProvider {
	return &CodexCLIProvider{
		command: config.Command, model: config.Model, timeout: config.Timeout.Duration,
		workspace: workspace,
	}
}

func (p *CodexCLIProvider) Name() string { return "codex-cli" }

func (p *CodexCLIProvider) Status(ctx context.Context) AIProviderStatus {
	p.mu.Lock()
	if time.Since(p.statusAt) < 30*time.Second {
		status := p.status
		p.mu.Unlock()
		return status
	}
	p.mu.Unlock()
	status := p.probe(ctx)
	p.mu.Lock()
	p.status, p.statusAt = status, time.Now()
	p.mu.Unlock()
	return status
}

func (p *CodexCLIProvider) probe(ctx context.Context) AIProviderStatus {
	command, err := p.secureCommand()
	if err != nil {
		return AIProviderStatus{Summary: "unavailable", Detail: err.Error()}
	}
	probeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	output, err := exec.CommandContext(probeCtx, command, "login", "status").CombinedOutput()
	if err != nil || !strings.Contains(strings.ToLower(string(output)), "logged in") {
		return AIProviderStatus{Summary: "not authenticated", Detail: safeAIText(string(output))}
	}
	model := p.model
	if model == "" {
		model, err = bundledDefaultCodexModel(probeCtx, command)
		if err != nil {
			return AIProviderStatus{
				Summary: "model discovery failed", Detail: err.Error(),
			}
		}
	}
	return AIProviderStatus{Ready: true, Summary: "ready via ChatGPT login", Detail: command, Model: model}
}

func bundledDefaultCodexModel(ctx context.Context, command string) (string, error) {
	output, err := exec.CommandContext(ctx, command, "debug", "models", "--bundled").Output()
	if err != nil {
		return "", fmt.Errorf("query bundled Codex models: %w", err)
	}
	var catalog struct {
		Models []struct {
			Slug       string `json:"slug"`
			Visibility string `json:"visibility"`
			Priority   int    `json:"priority"`
		} `json:"models"`
	}
	if err := json.Unmarshal(output, &catalog); err != nil {
		return "", fmt.Errorf("decode bundled Codex models: %w", err)
	}
	best := ""
	bestPriority := int(^uint(0) >> 1)
	for _, model := range catalog.Models {
		if model.Slug != "" && model.Visibility == "list" && model.Priority < bestPriority {
			best, bestPriority = model.Slug, model.Priority
		}
	}
	if best == "" {
		return "", fmt.Errorf("Codex reported no selectable bundled model")
	}
	return best, nil
}

func (p *CodexCLIProvider) secureCommand() (string, error) {
	command, err := exec.LookPath(p.command)
	if err != nil {
		return "", fmt.Errorf("find Codex CLI: %w", err)
	}
	command, err = filepath.EvalSymlinks(command)
	if err != nil {
		return "", fmt.Errorf("resolve Codex CLI: %w", err)
	}
	info, err := os.Stat(command)
	if err != nil {
		return "", err
	}
	if info.Mode().Perm()&0o022 != 0 {
		return "", fmt.Errorf("refusing group/world-writable Codex executable %s", command)
	}
	workspace, _ := filepath.Abs(p.workspace)
	if pathWithin(workspace, command) {
		return "", fmt.Errorf("refusing Codex executable inside the reviewed workspace")
	}
	for _, name := range []string{"HOME", "CODEX_HOME"} {
		if value := os.Getenv(name); value != "" && pathWithin(workspace, value) {
			return "", fmt.Errorf("refusing %s inside the reviewed workspace", name)
		}
	}
	return command, nil
}

func pathWithin(parent, candidate string) bool {
	parent, parentErr := filepath.Abs(parent)
	candidate, candidateErr := filepath.Abs(candidate)
	if parentErr != nil || candidateErr != nil {
		return false
	}
	rel, err := filepath.Rel(parent, candidate)
	return err == nil && rel != ".." &&
		!strings.HasPrefix(rel, ".."+string(filepath.Separator))
}

func (p *CodexCLIProvider) Generate(ctx context.Context, request AIInferenceRequest) (AIInferenceResponse, error) {
	return p.generate(ctx, request, nil)
}

func (p *CodexCLIProvider) GenerateWithProgress(
	ctx context.Context,
	request AIInferenceRequest,
	report func(AIProviderProgress),
) (AIInferenceResponse, error) {
	return p.generate(ctx, request, report)
}

func (p *CodexCLIProvider) generate(
	ctx context.Context,
	request AIInferenceRequest,
	report func(AIProviderProgress),
) (AIInferenceResponse, error) {
	command, err := p.secureCommand()
	if err != nil {
		return AIInferenceResponse{}, err
	}
	temp, err := os.MkdirTemp("", "diple-ai-*")
	if err != nil {
		return AIInferenceResponse{}, err
	}
	defer os.RemoveAll(temp)
	if pathWithin(p.workspace, temp) {
		return AIInferenceResponse{}, fmt.Errorf("refusing AI temporary directory inside the reviewed workspace")
	}
	if err := os.Chmod(temp, 0o700); err != nil {
		return AIInferenceResponse{}, err
	}
	schemaPath := filepath.Join(temp, "response-schema.json")
	if err := os.WriteFile(schemaPath, request.Schema, 0o600); err != nil {
		return AIInferenceResponse{}, err
	}
	runCtx, cancel := context.WithTimeout(ctx, p.timeout)
	defer cancel()
	args := []string{
		"exec", "--ignore-user-config", "--ignore-rules", "--strict-config", "--ephemeral",
		"--skip-git-repo-check", "-C", temp, "--sandbox", "read-only",
		"-c", `approval_policy="never"`,
		"--disable", "shell_tool", "--disable", "unified_exec", "--disable", "code_mode",
		"--disable", "code_mode_host", "--disable", "shell_snapshot",
		"--disable", "apps", "--disable", "browser_use", "--disable", "browser_use_external",
		"--disable", "browser_use_full_cdp_access", "--disable", "in_app_browser",
		"--disable", "standalone_web_search", "--disable", "computer_use",
		"--disable", "image_generation", "--disable", "plugins", "--disable", "skill_search",
		"--disable", "skill_mcp_dependency_install", "--disable", "memories",
		"--disable", "multi_agent", "--disable", "multi_agent_v2",
		"--disable", "auth_elicitation", "--disable", "tool_call_mcp_elicitation",
		"--disable", "request_permissions_tool", "--disable", "tool_suggest",
		"--disable", "hooks", "--disable", "remote_plugin", "--disable", "network_proxy",
		"--disable", "workspace_dependencies", "--disable", "goals",
		"--output-schema", schemaPath, "--json",
	}
	if request.Model != "" {
		args = append(args, "-m", request.Model)
	}
	args = append(args, "-")
	cmd := exec.CommandContext(runCtx, command, args...)
	cmd.Dir = temp
	cmd.Env = codexSafeEnvironment()
	cmd.Stdin = strings.NewReader(request.Prompt)
	var stdout, stderr limitedBuffer
	stdout.limit, stderr.limit = 4<<20, 64<<10
	progressWriter := &codexProgressWriter{output: &stdout, report: report}
	cmd.Stdout, cmd.Stderr = progressWriter, &stderr
	if err := cmd.Run(); err != nil {
		progressWriter.Flush()
		if runCtx.Err() != nil {
			return AIInferenceResponse{}, fmt.Errorf("Codex CLI: %w", runCtx.Err())
		}
		detail := codexFailureDetail(stdout.buffer.Bytes(), stderr.buffer.Bytes())
		return AIInferenceResponse{}, fmt.Errorf("Codex CLI: %w: %s", err, detail)
	}
	progressWriter.Flush()
	content, model, err := parseCodexEvents(stdout.buffer.Bytes())
	if err != nil {
		return AIInferenceResponse{}, err
	}
	if model == "" {
		model = request.Model
	}
	return AIInferenceResponse{Model: model, Content: content}, nil
}

type codexProgressWriter struct {
	output  *limitedBuffer
	report  func(AIProviderProgress)
	pending []byte
}

func (w *codexProgressWriter) Write(value []byte) (int, error) {
	n, err := w.output.Write(value)
	if n > 0 {
		w.pending = append(w.pending, value[:n]...)
		w.consume(false)
	}
	return n, err
}

func (w *codexProgressWriter) Flush() {
	w.consume(true)
}

func (w *codexProgressWriter) consume(flush bool) {
	for {
		index := bytes.IndexByte(w.pending, '\n')
		if index < 0 {
			if flush && len(w.pending) > 0 {
				w.reportLine(w.pending)
				w.pending = nil
			}
			return
		}
		w.reportLine(w.pending[:index])
		w.pending = w.pending[index+1:]
	}
}

func (w *codexProgressWriter) reportLine(line []byte) {
	if w.report == nil || len(bytes.TrimSpace(line)) == 0 {
		return
	}
	var event map[string]any
	if json.Unmarshal(line, &event) != nil {
		return
	}
	eventType, _ := event["type"].(string)
	if eventType == "turn.started" {
		w.report(AIProviderProgress{Kind: "activity", Text: "Model started"})
		return
	}
	item, _ := event["item"].(map[string]any)
	itemType, _ := item["type"].(string)
	if itemType != "reasoning" {
		return
	}
	text, _ := item["text"].(string)
	text = truncateAIInline(safeAIText(text), 600)
	if text != "" {
		w.report(AIProviderProgress{Kind: "reasoning", Text: text})
	}
}

func codexFailureDetail(stdout, stderr []byte) string {
	var messages []string
	scanner := bufio.NewScanner(bytes.NewReader(stdout))
	scanner.Buffer(make([]byte, 4096), 4<<20)
	for scanner.Scan() {
		var event map[string]any
		if json.Unmarshal(scanner.Bytes(), &event) != nil {
			continue
		}
		eventType, _ := event["type"].(string)
		if eventType != "error" && eventType != "turn.failed" {
			continue
		}
		if message, ok := event["message"].(string); ok && strings.TrimSpace(message) != "" {
			messages = append(messages, message)
		}
		if failure, ok := event["error"].(map[string]any); ok {
			if message, ok := failure["message"].(string); ok && strings.TrimSpace(message) != "" {
				messages = append(messages, message)
			}
		}
	}
	if len(messages) > 0 {
		return truncateAIInline(safeAIText(strings.Join(messages, "; ")), 2_000)
	}
	if detail := safeAIText(string(stderr)); detail != "" {
		return truncateAIInline(detail, 2_000)
	}
	return "Codex returned no diagnostic output"
}

func codexSafeEnvironment() []string {
	allowed := map[string]bool{
		"HOME": true, "CODEX_HOME": true, "PATH": true,
		"SSL_CERT_FILE": true, "SSL_CERT_DIR": true,
		"TERM": true,
	}
	var result []string
	for _, item := range os.Environ() {
		name := strings.SplitN(item, "=", 2)[0]
		if allowed[name] {
			result = append(result, item)
		}
	}
	return result
}

func parseCodexEvents(data []byte) ([]byte, string, error) {
	scanner := bufio.NewScanner(bytes.NewReader(data))
	scanner.Buffer(make([]byte, 4096), 4<<20)
	var content []byte
	var model string
	for scanner.Scan() {
		var event map[string]any
		if err := json.Unmarshal(scanner.Bytes(), &event); err != nil {
			return nil, "", fmt.Errorf("invalid Codex event stream: %w", err)
		}
		eventType, _ := event["type"].(string)
		lower := strings.ToLower(eventType)
		if strings.Contains(lower, "tool") || strings.Contains(lower, "command") ||
			strings.Contains(lower, "file_change") {
			return nil, "", fmt.Errorf("Codex attempted forbidden capability %q", eventType)
		}
		if value, ok := event["model"].(string); ok {
			model = value
		}
		item, _ := event["item"].(map[string]any)
		itemType, _ := item["type"].(string)
		itemLower := strings.ToLower(itemType)
		if strings.Contains(itemLower, "tool") || strings.Contains(itemLower, "command") ||
			strings.Contains(itemLower, "file_change") {
			return nil, "", fmt.Errorf("Codex attempted forbidden item %q", itemType)
		}
		switch itemType {
		case "", "reasoning", "agent_message":
		default:
			return nil, "", fmt.Errorf("Codex emitted unsupported item %q", itemType)
		}
		if itemType == "agent_message" {
			if text, ok := item["text"].(string); ok {
				content = []byte(text)
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, "", err
	}
	if len(content) == 0 {
		return nil, "", fmt.Errorf("Codex returned no structured response")
	}
	return content, model, nil
}