ollie-core/agent/prompt_resolver.go

package agent

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"ollie/paths"
)

// PromptEnv returns the standard PRIME_* environment variables for prompt
// template resolution. For local sessions, values are detected from the
// filesystem. For remote sessions, callers should override these with
// values from remote.HostInfo.
func PromptEnv(cwd string) []string {
	platform := runtime.GOOS
	isGitRepo := "false"
	if cwd != "" {
		if _, err := os.Stat(filepath.Join(cwd, ".git")); err == nil {
			isGitRepo = "true"
		}
	}
	return []string{
		"PRIME_CWD=" + cwd,
		"PRIME_PLATFORM=" + platform,
		"PRIME_DATE=" + time.Now().Format("2006-01-02"),
		"PRIME_IS_GIT_REPO=" + isGitRepo,
	}
}

// resolvePrompt interprets a Prompt from an agent config.
//
// If the prompt was parsed from a JSON array, each element is executed as a
// shell command and the combined stdout (joined by newlines) is returned.
//
// If the prompt was parsed from a JSON string, the existing single-string
// semantics apply: environment variables are expanded, then:
//   - If the string contains a newline, it is treated as literal text.
//   - If the string starts with '!', the rest is executed as a shell command.
//   - If the expanded string names an existing file, the file is read.
//   - Otherwise the string is used as-is.
func resolvePrompt(p Prompt, cwd string, env []string) (string, error) {
	if len(p.Value) == 0 {
		return "", nil
	}
	if p.IsExec {
		return resolveExecPrompt(p.Value, cwd, env)
	}
	return resolveStringPrompt(p.Value[0], cwd, env)
}

func resolveExecPrompt(cmds []string, cwd string, env []string) (string, error) {
	// Build env map for variable expansion.
	fullEnv := append(os.Environ(), env...)
	envMap := make(map[string]string, len(fullEnv))
	for _, e := range fullEnv {
		if k, v, ok := strings.Cut(e, "="); ok {
			envMap[k] = v
		}
	}
	expand := func(s string) string {
		return os.Expand(s, func(key string) string {
			return envMap[key]
		})
	}

	var parts []string
	for _, entry := range cmds {
		entry = strings.TrimSpace(entry)
		if entry == "" {
			continue
		}
		// Try resolving as a prompt file first.
		if content, ok := resolvePromptFile(entry, envMap, expand); ok {
			if content != "" {
				parts = append(parts, content)
			}
			continue
		}
		// Fall back to shell execution.
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		cmd := exec.CommandContext(ctx, "sh", "-c", entry)
		if cwd != "" {
			if info, err := os.Stat(cwd); err == nil && info.IsDir() {
				cmd.Dir = cwd
			}
		}
		cmd.Env = fullEnv
		var stdout, stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		err := cmd.Run()
		cancel()
		if err != nil {
			return "", fmt.Errorf("prompt command %q failed: %v: %s", entry, err, stderr.String())
		}
		if out := strings.TrimRight(stdout.String(), "\n"); out != "" {
			parts = append(parts, out)
		}
	}
	return strings.Join(parts, "\n"), nil
}

// resolvePromptFile attempts to resolve entry as a prompt file.
// It first checks if the entry (after env expansion) is an absolute path to an
// existing file. Otherwise it searches CfgDir()/prompts for entry and entry.md.
// Returns the file content (with env vars expanded) and true if found.
func resolvePromptFile(entry string, envMap map[string]string, expand func(string) string) (string, bool) {
	expanded := expand(entry)

	// If it contains shell metacharacters, it's a command, not a file name.
	if strings.ContainsAny(expanded, "|&;<>()$`\\\"' \t") {
		return "", false
	}

	// Try as absolute path.
	if filepath.IsAbs(expanded) {
		if info, err := os.Stat(expanded); err == nil && !info.IsDir() {
			return readAndExpand(expanded, expand)
		}
		// Try with .md suffix.
		if info, err := os.Stat(expanded + ".md"); err == nil && !info.IsDir() {
			return readAndExpand(expanded+".md", expand)
		}
		return "", false
	}

	// Search CfgDir()/prompts.
	for _, dir := range paths.PromptsDirs() {
		if dir == "" {
			continue
		}
		for _, candidate := range []string{
			filepath.Join(dir, expanded),
			filepath.Join(dir, expanded+".md"),
		} {
			if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
				return readAndExpand(candidate, expand)
			}
		}
	}
	return "", false
}

func readAndExpand(path string, expand func(string) string) (string, bool) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", false
	}
	content := strings.TrimRight(expand(string(data)), "\n")
	return content, true
}

func resolveStringPrompt(prompt, cwd string, env []string) (string, error) {
	if prompt == "" {
		return "", nil
	}
	envMap := make(map[string]string, len(env))
	for _, e := range env {
		if k, v, ok := strings.Cut(e, "="); ok {
			envMap[k] = v
		}
	}
	expanded := os.Expand(prompt, func(key string) string {
		if val, ok := envMap[key]; ok {
			return val
		}
		if val := os.Getenv(key); val != "" {
			return val
		}
		return ""
	})
	expanded = strings.TrimSpace(expanded)
	if expanded == "" {
		return "", nil
	}
	if strings.Contains(expanded, "\n") {
		return expanded, nil
	}
	if strings.HasPrefix(expanded, "!") {
		cmdStr := strings.TrimPrefix(expanded, "!")
		cmdStr = strings.TrimSpace(cmdStr)
		if cmdStr == "" {
			return "", nil
		}
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		cmd := exec.CommandContext(ctx, "sh", "-c", cmdStr)
		var cmdEnv []string
		if cwd != "" {
			if info, err := os.Stat(cwd); err == nil && info.IsDir() {
				cmd.Dir = cwd
			} else {
				cmdEnv = append(cmdEnv, "PWD="+cwd)
			}
		}
		if len(cmdEnv) > 0 {
			cmd.Env = append(os.Environ(), cmdEnv...)
		}
		var stdout, stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			return "", fmt.Errorf("prompt command failed: %v: %s", err, stderr.String())
		}
		return strings.TrimRight(stdout.String(), "\n"), nil
	}
	path := expanded
	if !filepath.IsAbs(path) && cwd != "" {
		path = filepath.Join(cwd, path)
	}
	if info, err := os.Stat(path); err == nil && !info.IsDir() {
		data, err := os.ReadFile(path)
		if err != nil {
			return "", fmt.Errorf("read prompt file: %w", err)
		}
		return strings.TrimRight(string(data), "\n"), nil
	}
	return expanded, nil
}