BlackBean/pigo/internal/runtime/headless.go

// This file implements the headless / stdio run modes (US-020, FR-18): a
// non-interactive driver that runs the agent loop over a single prompt for
// scripting and CI. Two output modes are supported, mirroring pi's print-mode
// and rpc/stream-json protocols:
//
//   - PrintMode: run the loop to completion and write only the final assistant
//     text to the output (the "-p / --print" mode).
//   - StreamJSONMode: serialize every AgentEvent as a line-delimited JSON object
//     as it is emitted (the "--output-format stream-json" mode), so a parent
//     process can consume the run incrementally.
//
// The run's success/failure is reported as a returned error so the CLI can map
// it to a process exit code: a run whose final assistant message carries
// stopReason error/aborted, or whose stream result errors, is a failure.
package runtime

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/smallnest/pigo/internal/agentcore"
)

// HeadlessMode selects how a headless run reports its progress and result.
type HeadlessMode int

const (
	// PrintMode runs the loop to completion and writes only the final assistant
	// text to the output writer.
	PrintMode HeadlessMode = iota
	// StreamJSONMode writes each AgentEvent as a line-delimited JSON object as it
	// is emitted.
	StreamJSONMode
)

// HeadlessConfig configures a headless run.
type HeadlessConfig struct {
	// Run is the loop configuration (provider stream, tools, hooks).
	Run RunConfig
	// Mode selects print vs stream-json output. Defaults to PrintMode.
	Mode HeadlessMode
	// Out receives the run output (final text or JSON lines). Required.
	Out io.Writer
	// OnEvent, when non-nil, is invoked for every AgentEvent before output
	// handling. It is the seam plugin lifecycle-event delivery (US-017, #133)
	// hooks into, independent of the output mode. It must not block.
	OnEvent func(ev agentcore.AgentEvent)
	// Progress receives human-readable sub-agent progress lines
	// (SubAgentProgressEvent). Defaults to os.Stderr when nil. Progress is
	// stderr-only by contract: it is never serialised onto Out (stdout), so it
	// cannot pollute the final result text or the stream-json envelope stream.
	Progress io.Writer
}

// ErrRunFailed is the sentinel returned by RunHeadless when the agent run ended
// in a failure state (stopReason error/aborted). The CLI maps a non-nil error
// to a non-zero exit code.
type ErrRunFailed struct {
	// Reason is the stopReason (or message) that marked the run as failed.
	Reason string
}

func (e *ErrRunFailed) Error() string {
	if e.Reason == "" {
		return "agent run failed"
	}
	return "agent run failed: " + e.Reason
}

// RunHeadless runs the agent loop for the already-assembled agentCtx and drives
// output per cfg.Mode. It blocks until the run ends and returns nil on success
// or an error describing the failure (for exit-code mapping). It never returns
// before the stream is fully drained, so no goroutine is leaked.
func RunHeadless(ctx context.Context, agentCtx *agentcore.AgentContext, cfg HeadlessConfig) error {
	if cfg.Out == nil {
		return fmt.Errorf("headless: nil output writer")
	}
	stream := agentLoop(ctx, agentCtx, cfg.Run)

	// writeErr holds the first stream-json write failure. We keep draining after
	// it (DrainStream never returns early) so the loop's producer goroutine never
	// blocks on a synchronous Emit — honoring the no-leak contract on a broken
	// pipe. In stream-json mode every event is serialised; print mode only needs
	// the final message, which DrainStream returns.
	var writeErr error
	h := StreamHandler{}
	// Compose the output-mode serialiser (stream-json only) with the optional
	// external OnEvent (plugin lifecycle delivery, US-017). Both observe every
	// event; the serialiser runs first so a write failure is recorded even when a
	// plugin observer is also wired.
	//
	// SubAgentProgressEvent (D-9) is special-cased: it is written as a
	// human-readable line to the progress writer (stderr) and is deliberately
	// excluded from the stream-json stdout path so it never pollutes the result
	// output or the machine-readable envelope stream (progress is stderr-only).
	progress := cfg.Progress
	if progress == nil {
		progress = os.Stderr
	}
	streamJSON := cfg.Mode == StreamJSONMode
	h.OnEvent = func(ev agentcore.AgentEvent) {
		if pe, ok := ev.(agentcore.SubAgentProgressEvent); ok {
			writeProgressLine(progress, pe)
			if cfg.OnEvent != nil {
				cfg.OnEvent(ev)
			}
			return
		}
		if streamJSON && writeErr == nil {
			if err := writeEventJSON(cfg.Out, ev); err != nil {
				writeErr = err
			}
		}
		if cfg.OnEvent != nil {
			cfg.OnEvent(ev)
		}
	}
	lastAssistant, resErr := DrainStream(ctx, stream, h)
	if writeErr != nil {
		return writeErr
	}
	if resErr != nil {
		return resErr
	}

	if cfg.Mode == PrintMode {
		text := ""
		if lastAssistant != nil {
			text = agentcore.ContentToText(lastAssistant.Content)
		}
		if _, err := io.WriteString(cfg.Out, text); err != nil {
			return err
		}
		if text != "" && !strings.HasSuffix(text, "\n") {
			if _, err := io.WriteString(cfg.Out, "\n"); err != nil {
				return err
			}
		}
	}

	if lastAssistant != nil {
		switch lastAssistant.StopReason {
		case agentcore.StopReasonError:
			reason := lastAssistant.ErrorMessage
			if reason == "" {
				reason = "error"
			}
			return &ErrRunFailed{Reason: reason}
		case agentcore.StopReasonAborted:
			return &ErrRunFailed{Reason: "aborted"}
		}
	}
	return nil
}

// writeProgressLine renders one SubAgentProgressEvent as a human-readable line
// on w (stderr by contract). When the task supplied a description it is shown
// alongside the activity; otherwise the line degrades to the activity alone
// (Description MAY be empty, Activity never is). Write errors are ignored:
// progress is a non-critical, best-effort side channel.
func writeProgressLine(w io.Writer, ev agentcore.SubAgentProgressEvent) {
	if ev.Description != "" {
		fmt.Fprintf(w, "  ⏺ %s · %s\n", ev.Description, ev.Activity)
		return
	}
	fmt.Fprintf(w, "  ⏺ %s\n", ev.Activity)
}

// writeEventJSON serializes one AgentEvent as a single line of JSON, terminated
// by a newline, onto w. The envelope always carries a "type" discriminant so a
// consumer can dispatch without positional knowledge.
func writeEventJSON(w io.Writer, ev agentcore.AgentEvent) error {
	env := eventEnvelope(ev)
	b, err := json.Marshal(env)
	if err != nil {
		return fmt.Errorf("headless: marshal event: %w", err)
	}
	b = append(b, '\n')
	_, err = w.Write(b)
	return err
}

// eventEnvelope maps an AgentEvent onto a JSON-serializable object with a
// "type" discriminant plus the event's observable payload. Only fields that are
// safe and useful over the wire are included (assistant text, tool ids/names,
// stop reasons) — never secrets.
func eventEnvelope(ev agentcore.AgentEvent) map[string]any {
	env := map[string]any{"type": ev.EventType()}
	switch e := ev.(type) {
	case agentcore.AgentStartEvent:
		// The first event carries the backing session id (mirrors pi/Claude Code),
		// so a consumer can associate the run's output with a session and resume
		// it later. Omitted only when the run has no backing session (SessionID
		// unset), which the envelope treats as "not resumable".
		if e.SessionID != "" {
			env["sessionId"] = e.SessionID
		}
	case agentcore.AgentEndEvent:
		env["messageCount"] = len(e.Messages)
	case agentcore.TurnEndEvent:
		env["stopReason"] = e.Message.StopReason
		if text := agentcore.ContentToText(e.Message.Content); text != "" {
			env["text"] = text
		}
		if calls := e.Message.ToolCalls(); len(calls) > 0 {
			names := make([]string, len(calls))
			for i, c := range calls {
				names[i] = c.Name
			}
			env["toolCalls"] = names
		}
	case agentcore.MessageUpdateEvent:
		if a, ok := e.Message.(agentcore.AssistantMessage); ok {
			if text := agentcore.ContentToText(a.Content); text != "" {
				env["text"] = text
			}
		}
	case agentcore.ToolExecutionStartEvent:
		env["toolCallId"] = e.ToolCallID
		env["toolName"] = e.ToolName
	case agentcore.ToolExecutionEndEvent:
		env["toolCallId"] = e.ToolCallID
		env["toolName"] = e.ToolName
		env["isError"] = e.IsError
	case agentcore.CompactionStartEvent:
		env["reason"] = e.Reason
		env["tokensBefore"] = e.TokensBefore
	case agentcore.CompactionEvent:
		env["reason"] = e.Reason
		env["tokensBefore"] = e.TokensBefore
		env["tokensAfter"] = e.TokensAfter
		env["summarizedCount"] = e.SummarizedCount
		env["keptCount"] = e.KeptCount
		if e.ErrorMessage != "" {
			env["error"] = e.ErrorMessage
		}
	case agentcore.TelemetryEvent:
		// The run-end telemetry summary: structured metrics a script can read
		// directly from the stream-json output (observability -- structured telemetry collection). Per-tool
		// timings are flattened into a name→{count,totalMs} object so a JSON
		// consumer can index by tool name.
		env["turns"] = e.Turns
		env["truncationCount"] = e.TruncationCount
		env["compactionCount"] = e.CompactionCount
		env["contextUtilization"] = e.ContextUtilization
		env["contextTokens"] = e.ContextTokens
		env["contextWindow"] = e.ContextWindow
		tools := make(map[string]map[string]any, len(e.ToolDurationsMs))
		for name, t := range e.ToolDurationsMs {
			tools[name] = map[string]any{"count": t.Count, "totalMs": t.TotalMs}
		}
		env["toolDurationsMs"] = tools
	}
	return env
}