docker-agent/pkg/evaluation/eval.go

// Package evaluation provides an evaluation framework for testing agents.
package evaluation

import (
	"bufio"
	"bytes"
	"cmp"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"time"
	"uuid"

	"golang.org/x/sync/singleflight"

	"github.com/docker/docker-agent/pkg/config"
	"github.com/docker/docker-agent/pkg/config/latest"
	"github.com/docker/docker-agent/pkg/config/sources"
	"github.com/docker/docker-agent/pkg/session"
)

// Runner runs evaluations against an agent.
type Runner struct {
	Config

	agentSource config.Source
	agentConfig *latest.Config
	judge       *Judge
	runConfig   *config.RuntimeConfig

	// imageCache caches built container images by (workingDir, image) pair.
	imageCache   map[imageKey]string
	imageCacheMu sync.Mutex

	// imageBuildGroup deduplicates concurrent image builds for the same (workingDir, image) pair.
	imageBuildGroup singleflight.Group
}

// newRunner creates a new evaluation runner.
func newRunner(agentSource config.Source, runConfig *config.RuntimeConfig, cfg Config) *Runner {
	if runConfig != nil && runConfig.EncryptedConfig == "" {
		if source, ok := agentSource.(config.EncryptedConfigSource); ok && source.EncryptedConfig() != "" {
			runConfig = runConfig.Clone()
			runConfig.EncryptedConfig = source.EncryptedConfig()
		}
	}
	return &Runner{
		Config:      cfg,
		agentSource: agentSource,
		runConfig:   runConfig,
		imageCache:  make(map[imageKey]string),
	}
}

// Evaluate runs evaluations with a specified run name.
// ttyOut is used for progress bar rendering (should be the console/TTY).
// out is used for results and status messages (can be tee'd to a log file).
func Evaluate(ctx context.Context, ttyOut, out io.Writer, isTTY bool, runName string, runConfig *config.RuntimeConfig, cfg Config) (*EvalRun, error) {
	agentSource, err := sources.Resolve(cfg.AgentFilename, nil)
	if err != nil {
		return nil, fmt.Errorf("resolving agent: %w", err)
	}
	agentConfig, err := config.Load(ctx, agentSource, config.WithFlavors(runConfig.Flavors...))
	if err != nil {
		return nil, fmt.Errorf("loading agent: %w", err)
	}

	runner := newRunner(agentSource, runConfig, cfg)
	runner.agentConfig = agentConfig

	fmt.Fprintf(out, "Evaluation run: %s\n", runName)

	startTime := time.Now()
	results, err := runner.Run(ctx, ttyOut, out, isTTY)
	duration := time.Since(startTime)
	cfg.Flavors = slices.Clone(runConfig.Flavors)

	summary := computeSummary(results)
	summary.RepeatMetrics = computeRepeatMetrics(results, cfg.Repeat)
	printSummary(out, summary, duration)

	run := &EvalRun{
		Name:      runName,
		Timestamp: startTime,
		Duration:  duration,
		Config:    cfg,
		Results:   results,
		Summary:   summary,
	}

	if err != nil {
		return run, fmt.Errorf("running evaluations: %w", err)
	}

	return run, nil
}

// workItem represents a single evaluation to be processed.
type workItem struct {
	index int
	eval  *InputSession
}

// Run executes all evaluations concurrently and returns results.
// ttyOut is used for progress bar rendering (should be the console/TTY).
// out is used for results and status messages (can be tee'd to a log file).
func (r *Runner) Run(ctx context.Context, ttyOut, out io.Writer, isTTY bool) ([]Result, error) {
	fmt.Fprintln(out, "Loading evaluation sessions...")
	evals, err := r.loadEvalSessions(ctx)
	if err != nil {
		return nil, fmt.Errorf("loading evaluations: %w", err)
	}

	// Check whether any evaluations require relevance checking.
	// If so, the judge must be configured and working; validate eagerly
	// to fail fast on configuration issues (bad API key, wrong model, etc.)
	// instead of silently producing zero-relevance results.
	if needsJudge(evals) {
		if r.judge == nil {
			judge, err := createJudge(ctx, r.Config, r.runConfig, r.agentConfig)
			if err != nil {
				return nil, err
			}
			if judge == nil {
				return nil, errors.New("some evaluations have relevance criteria but no judge model is configured (use --judge-model)")
			}
			r.judge = judge
		}
		fmt.Fprintln(out, "Validating judge model...")
		if err := r.judge.Validate(ctx); err != nil {
			return nil, fmt.Errorf("%w", err)
		}
	}

	// Pre-build all unique container images in parallel before running evaluations.
	// This avoids serialized builds when multiple workers need the same image.
	if err := r.preBuildImages(ctx, out, evals); err != nil {
		return nil, fmt.Errorf("pre-building images: %w", err)
	}

	fmt.Fprintf(out, "Running %d evaluations with concurrency %d\n\n", len(evals), r.Concurrency)

	progress := newProgressBar(ttyOut, out, r.TTYFd, len(evals), isTTY)
	// Keep rendering until workers finish, even after evaluation cancellation.
	progressCtx, stopProgress := context.WithCancel(context.WithoutCancel(ctx))
	var progressWG sync.WaitGroup
	progressWG.Go(func() { progress.run(progressCtx) })
	defer func() {
		stopProgress()
		progressWG.Wait()
	}()

	results := make([]Result, len(evals))

	work := make(chan workItem, len(evals))
	for i := range evals {
		work <- workItem{index: i, eval: &evals[i]}
	}
	close(work)

	var wg sync.WaitGroup
	for range r.Concurrency {
		wg.Go(func() {
			for item := range work {
				if ctx.Err() != nil {
					return
				}

				progress.setRunning(item.eval.displayTitle())
				result, runErr := r.runSingleEval(ctx, item.eval)
				if runErr != nil {
					result.Error = runErr.Error()
					slog.ErrorContext(ctx, "Evaluation failed", "title", item.eval.displayTitle(), "error", runErr)
				}

				results[item.index] = result
				_, failures := result.checkResults()
				progress.complete(result.Title, len(failures) == 0)
				progress.printResult(result)
			}
		})
	}

	wg.Wait()

	if ctx.Err() != nil {
		return results, ctx.Err()
	}

	return results, nil
}

func (r *Runner) loadEvalSessions(ctx context.Context) ([]InputSession, error) {
	entries, err := os.ReadDir(r.EvalsDir)
	if err != nil {
		return nil, err
	}

	var evals []InputSession
	for _, entry := range entries {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		// Filter by --only patterns against file name if specified
		fileName := entry.Name()
		if len(r.Only) > 0 && !matchesAnyPattern(fileName, r.Only) {
			continue
		}

		if entry.IsDir() || !strings.HasSuffix(fileName, ".json") {
			continue
		}

		data, err := os.ReadFile(filepath.Join(r.EvalsDir, fileName))
		if err != nil {
			return nil, err
		}

		var evalSess session.Session
		if err := json.Unmarshal(data, &evalSess); err != nil {
			return nil, err
		}

		evals = append(evals, InputSession{
			Session:    &evalSess,
			SourcePath: filepath.Join(r.EvalsDir, fileName),
		})
	}

	// Sort by duration (longest first) to avoid long tail
	slices.SortFunc(evals, func(a, b InputSession) int {
		return cmp.Compare(b.Duration(), a.Duration())
	})

	// Repeat evals if requested
	repeat := max(r.Repeat, 1)
	if repeat > 1 {
		original := evals
		evals = make([]InputSession, 0, len(original)*repeat)
		for i := range repeat {
			for _, e := range original {
				evals = append(evals, InputSession{
					Session:     e.Session,
					SourcePath:  e.SourcePath,
					RepeatIndex: i + 1,
				})
			}
		}
	}

	return evals, nil
}

// preBuildImages pre-builds all unique container images needed for the evaluations.
// Concurrent calls for the same (workingDir, image) pair are deduplicated by
// getOrBuildImage's singleflight, so we simply iterate over all evals.
func (r *Runner) preBuildImages(ctx context.Context, out io.Writer, evals []InputSession) error {
	if len(evals) == 0 {
		return nil
	}

	// Count unique images to report an accurate number.
	unique := make(map[imageKey]struct{})
	for _, eval := range evals {
		criteria := eval.criteria()
		unique[imageKey{workingDir: criteria.WorkingDir, image: criteria.Image}] = struct{}{}
	}

	fmt.Fprintf(out, "Pre-building %d container image(s)...\n", len(unique))

	type buildResult struct {
		title string
		err   error
	}

	work := make(chan InputSession, len(evals))
	for _, eval := range evals {
		work <- eval
	}
	close(work)

	results := make(chan buildResult, len(evals))

	buildWorkers := min(r.Concurrency, len(evals))
	var wg sync.WaitGroup
	for range buildWorkers {
		wg.Go(func() {
			for eval := range work {
				if ctx.Err() != nil {
					results <- buildResult{title: eval.Title, err: ctx.Err()}
					continue
				}

				criteria := eval.criteria()

				_, err := r.getOrBuildImage(ctx, criteria)
				results <- buildResult{title: eval.Title, err: err}
			}
		})
	}

	go func() {
		wg.Wait()
		close(results)
	}()

	var errs []error
	for result := range results {
		if result.err != nil {
			errs = append(errs, fmt.Errorf("building image for %q: %w", result.title, result.err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("failed to build %d image(s): %w", len(errs), errs[0])
	}

	return nil
}

func (r *Runner) runSingleEval(ctx context.Context, evalSess *InputSession) (Result, error) {
	startTime := time.Now()
	title := evalSess.displayTitle()

	slog.DebugContext(ctx, "Starting evaluation", "title", title)

	evals := evalSess.criteria()

	userMessages := getUserMessages(evalSess.Session)

	result := Result{
		InputPath:         evalSess.SourcePath,
		Title:             title,
		Question:          strings.Join(userMessages, "\n"),
		SizeExpected:      evals.Size,
		RelevanceExpected: float64(len(evals.Relevance)),
		AssertionsTotal:   len(evals.Assertions),
	}

	expectedToolCalls := extractToolCalls(evalSess.Messages)
	if len(expectedToolCalls) > 0 {
		result.ToolCallsExpected = 1.0
	}

	imageID, err := r.getOrBuildImage(ctx, evals)
	if err != nil {
		return result, fmt.Errorf("building eval image: %w", err)
	}

	events, err := r.runDockerAgentInContainer(ctx, imageID, userMessages, evals.Setup)
	if err != nil {
		return result, fmt.Errorf("running docker agent in container: %w", err)
	}

	response, cost, outputTokens, actualToolCalls := parseContainerEvents(events)

	result.Response = response
	result.Cost = cost
	result.OutputTokens = outputTokens
	result.Size = getResponseSize(result.Response)

	// Build session from events for database storage
	result.Session = SessionFromEvents(events, title, userMessages)
	result.Session.InputID = evalSess.InputID
	result.Session.Evals = evals

	// Re-apply the display title in case a session_title event overrode it.
	// This ensures repeated evals retain their '#N' suffix in stored sessions.
	result.Session.SetTitle(title)

	if len(expectedToolCalls) > 0 || len(actualToolCalls) > 0 {
		result.ToolCallsScore = toolCallF1Score(expectedToolCalls, actualToolCalls)
	}

	if r.judge != nil && len(evals.Relevance) > 0 {
		// Use transcript for relevance checking to preserve temporal ordering
		transcript := buildTranscript(events)
		results, err := r.judge.CheckRelevance(ctx, transcript, evals.Relevance)
		if err != nil {
			return result, fmt.Errorf("relevance check failed: %w", err)
		}
		var passed float64
		for _, rr := range results {
			if rr.Passed {
				passed++
			}
		}
		result.RelevancePassed = passed
		result.RelevanceResults = results
	}

	// Run code-based assertions against the agent output.
	if len(evals.Assertions) > 0 {
		assertionResults := runAssertions(evals.Assertions, response, cost, actualToolCalls)
		result.AssertionResults = assertionResults
		for _, ar := range assertionResults {
			if ar.Passed {
				result.AssertionsPassed++
			}
		}
	}

	slog.DebugContext(ctx, "Evaluation complete", "title", title, "duration", time.Since(startTime))
	return result, nil
}

func (r *Runner) runDockerAgentInContainer(ctx context.Context, imageID string, questions []string, setup string) ([]map[string]any, error) {
	agentDir := r.agentSource.ParentDir()
	agentFile := filepath.Base(r.agentSource.Name())
	containerID := uuid.NewV4()
	containerName := fmt.Sprintf("docker-agent-eval-%d", binary.BigEndian.Uint32(containerID[:4]))

	args := []string{
		"run",
		"--name", containerName,
		"--privileged",
		"--init",
	}
	if !r.KeepContainers {
		args = append(args, "--rm")
	}
	args = append(args,
		"-i",
		"-v", agentDir+":/configs:ro",
	)

	judgeConfigDir, err := r.stageJudgeProvider()
	if err != nil {
		return nil, err
	}
	if judgeConfigDir != "" {
		defer os.RemoveAll(judgeConfigDir)
		args = append(args, "-v", judgeConfigDir+":/judge-config:ro")
	}

	var env []string
	// addEnv forwards a variable to the container: "-e NAME" tells the runtime
	// CLI to pass it through, and NAME=VALUE sets it on the CLI process.
	addEnv := func(name, value string) {
		args = append(args, "-e", name)
		env = append(env, name+"="+value)
	}

	// Only forward environment variables explicitly requested via -e.
	// Format: KEY or KEY=VALUE
	for _, entry := range r.EnvVars {
		if key, val, hasValue := strings.Cut(entry, "="); hasValue && key != "" {
			addEnv(key, val)
		} else if val, ok := r.runConfig.EnvProvider().Get(ctx, entry); ok && entry != "" {
			addEnv(entry, val)
		}
	}

	// When a setup script is provided, mount it into the container and
	// override the entrypoint to run it before docker agent run --exec.
	// The default entrypoint is: /run.sh /docker-agent run --exec --yolo --json
	// /run.sh starts dockerd then exec's "$@".
	if setup != "" {
		setupID := uuid.NewV4()
		setupFile := filepath.Join(os.TempDir(), fmt.Sprintf("docker-agent-eval-setup-%d.sh", binary.BigEndian.Uint32(setupID[:4])))
		if err := os.WriteFile(setupFile, []byte(setup), 0o600); err != nil {
			return nil, fmt.Errorf("writing setup script: %w", err)
		}
		defer os.Remove(setupFile)

		args = append(args,
			"-v", setupFile+":/setup.sh:ro",
			"--entrypoint", "/run.sh",
		)
	}

	args = append(args, imageID)

	if setup != "" {
		// Run setup script, then docker agent run --exec with the original arguments.
		args = append(args, "sh", "-c", "sh /setup.sh && exec /docker-agent run --exec --yolo --json \"$@\"", "--")
	}
	if judgeConfigDir != "" {
		args = append(args, "--config-dir", "/judge-config")
	}
	for _, flavor := range r.runConfig.Flavors {
		args = append(args, "--flavor", flavor)
	}
	args = append(args, "--", "/configs/"+agentFile)
	args = append(args, questions...)

	containerRuntime := r.containerRuntimeOrDefault()
	cmd := exec.CommandContext(ctx, containerRuntime, args...)
	cmd.Env = append(os.Environ(), env...)
	// On cancellation send SIGINT instead of the default SIGKILL: the
	// Docker-compatible CLI proxies SIGINT to the container (SIGKILL is never
	// proxied and would leave the container running daemon-side). The
	// container's exit also triggers --rm removal. WaitDelay force-kills the
	// CLI if the container doesn't stop in time.
	cmd.Cancel = func() error {
		return cmd.Process.Signal(os.Interrupt)
	}
	cmd.WaitDelay = 10 * time.Second
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("creating stdout pipe: %w", err)
	}
	// Let exec.Cmd own the stderr copy: Wait blocks until its own copying
	// goroutine finishes (bounded by WaitDelay), so stderrBuf is safe to read
	// after Wait returns. A hand-rolled StderrPipe+goroutine has no such
	// guarantee and races with the reads below.
	var stderrBuf bytes.Buffer
	cmd.Stderr = &stderrBuf

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("starting %s run: %w", containerRuntime, err)
	}

	var events []map[string]any
	scanner := bufio.NewScanner(stdout)
	scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024)

	for scanner.Scan() {
		line := scanner.Text()
		if line == "" {
			continue
		}

		var event map[string]any
		if err := json.Unmarshal([]byte(line), &event); err != nil {
			slog.DebugContext(ctx, "Failed to parse JSON event", "line", line, "error", err)
			continue
		}
		events = append(events, event)
	}

	if err := scanner.Err(); err != nil {
		slog.WarnContext(ctx, "Error reading container output", "error", err)
	}

	waitErr := cmd.Wait()
	if waitErr != nil {
		slog.DebugContext(ctx, "Container exited with error", "stderr", stderrBuf.String(), "error", waitErr)
	}

	if len(events) == 0 {
		stderrStr := strings.TrimSpace(stderrBuf.String())
		if waitErr != nil {
			return nil, fmt.Errorf("container failed: %w (stderr: %s)", waitErr, stderrStr)
		}
		if stderrStr != "" {
			return nil, fmt.Errorf("no events received from container (stderr: %s)", stderrStr)
		}
		return nil, errors.New("no events received from container")
	}

	return events, nil
}

func parseContainerEvents(events []map[string]any) (response string, cost float64, outputTokens int64, toolCalls []string) {
	var responseBuf strings.Builder
	outputs := acceptedStructuredOutputs(events)
	for _, event := range events {
		eventType, _ := event["type"].(string)

		switch eventType {
		case "agent_choice":
			if content, ok := event["content"].(string); ok {
				responseBuf.WriteString(content)
			}
		case "tool_call":
			if tc, ok := event["tool_call"].(map[string]any); ok {
				if fn, ok := tc["function"].(map[string]any); ok {
					if name, ok := fn["name"].(string); ok {
						toolCalls = append(toolCalls, name)
					}
				}
			}
		case "token_usage":
			if usage, ok := event["usage"].(map[string]any); ok {
				if c, ok := usage["cost"].(float64); ok {
					cost = c
				}
				if snapshotOnly, _ := usage["snapshot_only"].(bool); !snapshotOnly {
					if tokens, ok := usage["output_tokens"].(float64); ok {
						outputTokens += int64(tokens)
					}
				}
			}
		}
	}

	response = responseBuf.String()
	if len(outputs) > 0 {
		// Structured results are terminal messages, not streamed text. Use
		// the last answer, including a later handoff or conversational turn.
		response = SessionFromEvents(events, "", nil).GetLastAssistantMessageContent()
	}
	return response, cost, outputTokens, toolCalls
}

// buildTranscript creates a chronological transcript of agent interactions.
// Unlike parseContainerEvents which only extracts text, this preserves the
// temporal sequence of events, enabling evaluation of criteria like
// "explains before executing" or "announces tool usage beforehand".
func buildTranscript(events []map[string]any) string {
	var transcript strings.Builder
	var pendingText strings.Builder
	var currentAgent string
	terminations := &terminationTracker{}

	flushText := func() {
		if pendingText.Len() == 0 {
			return
		}
		fmt.Fprintf(&transcript, "[Agent %s says]:\n%s\n\n", cmp.Or(currentAgent, "unknown"), pendingText.String())
		pendingText.Reset()
	}

	outputs := acceptedStructuredOutputs(events)
	for index, event := range events {
		switch event["type"] {
		case "agent_choice":
			if agentName, _ := event["agent_name"].(string); agentName != "" {
				currentAgent = agentName
			}
			if content, _ := event["content"].(string); content != "" {
				pendingText.WriteString(content)
			}

		case "tool_call":
			flushText()
			name, args := getToolCallInfo(event)
			if agentName, _ := event["agent_name"].(string); agentName != "" {
				currentAgent = agentName
			}
			fmt.Fprintf(&transcript, "[Agent %s calls tool %q with arguments: %s]\n\n", cmp.Or(currentAgent, "unknown"), name, args)

		case "tool_call_response":
			// The ToolCallResponseEvent has tool_definition at the top level, not
			// nested under "tool_call".
			var name string
			if td, ok := event["tool_definition"].(map[string]any); ok {
				name, _ = td["name"].(string)
			}
			response, _ := event["response"].(string)
			if len(response) > 500 {
				response = response[:500] + "...(truncated)"
			}
			fmt.Fprintf(&transcript, "[Tool %q returns: %s]\n\n", name, response)
			if output, ok := outputs[index]; ok {
				agentName, _ := event["agent_name"].(string)
				fmt.Fprintf(&transcript, "[Agent %s final answer]:\n%s\n\n", cmp.Or(agentName, "unknown"), output)
			}

		case "budget_exceeded":
			// A budget stop is a structured termination, never an error. The
			// marker lands at the event's chronological position, immediately
			// followed by the assistant stop message the event embeds under
			// stop_message. All message_added events are ignored: they carry
			// no JSON payload, and regular assistant turns are already
			// captured from agent_choice events.
			flushText()
			term, stop := terminations.observeBudgetExceeded(event)
			if term != nil {
				fmt.Fprintf(&transcript, "%s\n\n", terminationTranscriptMarker(term))
			}
			if stop != nil {
				fmt.Fprintf(&transcript, "[Agent %s says]:\n%s\n\n", cmp.Or(stop.AgentName, currentAgent, "unknown"), stop.Message.Content)
			}
		}
	}

	flushText()
	return transcript.String()
}

// getToolCallInfo extracts the tool name and arguments from an event.
func getToolCallInfo(event map[string]any) (name, args string) {
	tc, _ := event["tool_call"].(map[string]any)
	fn, _ := tc["function"].(map[string]any)
	name, _ = fn["name"].(string)
	args, _ = fn["arguments"].(string)
	return name, args
}

// matchesAnyPattern returns true if the name contains any of the patterns (case-insensitive).
func matchesAnyPattern(name string, patterns []string) bool {
	nameLower := strings.ToLower(name)
	return slices.ContainsFunc(patterns, func(pattern string) bool {
		return strings.Contains(nameLower, strings.ToLower(pattern))
	})
}

// needsJudge returns true if any evaluation session has relevance criteria,
// meaning a judge model is required to evaluate them.
func needsJudge(evals []InputSession) bool {
	return slices.ContainsFunc(evals, func(s InputSession) bool {
		return s.Evals != nil && len(s.Evals.Relevance) > 0
	})
}