first commit

pigo/internal/cli/headless/headless.go

@@ -0,0 +1,205 @@
+// This file is the headless run driver: the print / stream-json run path
+// (US-020) extracted from the CLI dispatch seam (#363). dispatch resolves the
+// output mode and the run environment, then hands off to Run, which wires the
+// session, prompt, thinking level, and provider credentials into a
+// runtime.HeadlessConfig and executes one run. Plugin slash commands and output
+// mode parsing live here because they are specific to the headless path.
+package headless
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+	"github.com/smallnest/pigo/internal/cli/run"
+	"github.com/smallnest/pigo/internal/cli/ui"
+	"github.com/smallnest/pigo/internal/plugin"
+	"github.com/smallnest/pigo/internal/provider"
+	"github.com/smallnest/pigo/internal/runtime"
+)
+
+// RunParams carries the resolved inputs for one headless run. Mode and Env are
+// resolved by the caller (dispatch) — Mode via ParseOutputMode, Env via
+// run.SetupEnv — so their distinct exit codes stay at the call site; Run owns
+// the rest of the run lifecycle.
+type RunParams struct {
+	Mode          runtime.HeadlessMode
+	Env           run.Env
+	Prompt        string
+	Model         string
+	APIKey        string
+	ThinkingLevel string
+	ResumeID      string
+}
+
+// Run executes one headless run over p.Prompt, writing agent output to out and
+// diagnostics to errOut, and returns a process exit code (0 = success). The run
+// is backed by a session so its id appears in the first stream-json event and it
+// can be resumed with --resume/--continue; a resumed session seeds its prior
+// messages ahead of the new prompt.
+func Run(ctx context.Context, p RunParams, out, errOut io.Writer) int {
+	env := p.Env
+	// Best-effort plugin slash-command support in headless mode: if the prompt is
+	// a "/cmd ..." naming a plugin command, invoke it, print its notifications to
+	// errOut, and use the returned prompt for this run (appending the raw args if
+	// the command produced no prompt). Headless has no turn injection, so
+	// appending the returned prompt is the accepted behavior. A non-plugin prompt
+	// or unknown command is left untouched.
+	headlessPrompt := resolveHeadlessPluginCommand(p.Prompt, env.Plugins, errOut)
+	promptContent, err := ui.BuildUserContent(headlessPrompt)
+	if err != nil {
+		fmt.Fprintf(errOut, "pigo: %v\n", err)
+		return 1
+	}
+
+	// Back the headless run with a session so its id appears in the first
+	// stream-json event and the run can be resumed with --resume/--continue,
+	// matching the interactive REPL and pi/Claude Code. A resumed session seeds
+	// its prior messages ahead of the new prompt.
+	priorMsgs, hs, err := openHeadlessSession(p.ResumeID, p.Model, env.ProviderName, env.SysPrompt)
+	if err != nil {
+		fmt.Fprintf(errOut, "pigo: %v\n", err)
+		return 1
+	}
+	messages := append(priorMsgs, agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: promptContent})
+	agentCtx := &agentcore.AgentContext{
+		SystemPrompt: hs.header.SystemPrompt,
+		Messages:     messages,
+		Tools:        env.Tools,
+	}
+
+	// Resolve the effective reasoning-effort level through the layered config
+	// chain (default < global < project < env < --thinking-level flag).
+	thinking, err := run.ResolveThinkingLevel(p.ThinkingLevel)
+	if err != nil {
+		fmt.Fprintf(errOut, "pigo: %v\n", err)
+		return 2
+	}
+
+	// Resolve the API key by provider name from the environment (never logged).
+	// An explicit --api-key overrides env/config for the resolved provider.
+	creds := provider.NewCredentialStore(nil)
+	creds.SetOverride(env.ProviderName, p.APIKey)
+	runCfg := run.NewConfig(p.Model, env.ProviderName, thinking, env.Provider, creds, run.ToolRegistry(env.Tools), run.TodoReminders(env.Tools))
+	runCfg.SessionID = hs.header.ID
+	// Route auto-compaction checkpoints to the shared memory root so a rebuild can
+	// recover the pre-watermark prefix (no-op when memory is disabled → empty root).
+	runCfg.MemoryRoot = run.MemoryRootFromTools(env.Tools)
+
+	// Wire hooks uniformly with every other driver (#425): resolve the trust-gated
+	// hook set, install the tool-execution + Stop seams, dispatch SessionStart, and
+	// chain the SessionEnd/PreCompact observer onto the plugin event notifier. A
+	// malformed hook layer is a config error (exit 2), matching thinking-level.
+	source := "startup"
+	if p.ResumeID != "" {
+		source = "resume"
+	}
+	set, herr := run.ResolveHookSet(env.Cwd, run.Trusted(env.Cwd))
+	if herr != nil {
+		fmt.Fprintf(errOut, "pigo: %v\n", herr)
+		return 2
+	}
+	hookDeps := run.HookDeps{SessionID: hs.header.ID, ProjectDir: env.Cwd, WarnLog: errOut}
+	// Deliver agent lifecycle events to any subscribed plugin (US-017, #133).
+	// NewEventNotifier returns nil when no plugin subscribes, so the base handler
+	// stays nil in the common no-plugin case.
+	var baseOnEvent func(agentcore.AgentEvent)
+	if n := plugin.NewEventNotifier(env.Plugins, errOut); n != nil {
+		baseOnEvent = n.Handle
+	}
+	d, onEvent := run.InstallDriverHooks(ctx, &runCfg, set, hookDeps, source, baseOnEvent)
+	// UserPromptSubmit runs before the prompt is handed to the loop: a block aborts
+	// the headless run non-zero; additionalContext is injected into this run only.
+	if d != nil {
+		if block, reason := run.DispatchUserPromptSubmit(ctx, d, &runCfg, hookDeps, headlessPrompt); block {
+			fmt.Fprintf(errOut, "pigo: prompt blocked by hook: %s\n", reason)
+			return 1
+		}
+	}
+
+	cfg := runtime.HeadlessConfig{
+		Mode: p.Mode,
+		Out:  out,
+		Run:  runCfg,
+	}
+	cfg.OnEvent = onEvent
+	runErr := runtime.RunHeadless(ctx, agentCtx, cfg)
+	// Persist the run's messages regardless of run outcome so a partial run is
+	// still resumable; a persistence failure is reported but does not mask a run
+	// error.
+	if perr := hs.persist(agentCtx); perr != nil {
+		fmt.Fprintf(errOut, "pigo: warning: could not persist session %s: %v\n", hs.header.ID, perr)
+	}
+	if runErr != nil {
+		fmt.Fprintf(errOut, "pigo: %v\n", runErr)
+		return 1
+	}
+	return 0
+}
+
+// resolveHeadlessPluginCommand gives the headless / print path best-effort
+// support for plugin slash commands. When prompt is a "/cmd ..." naming a
+// plugin command (from mgr.Commands()), it invokes the command, prints each
+// returned notification to notifyOut, and returns the command's returned Prompt
+// as the run's prompt. If the command returns no prompt, the raw argument text
+// is used instead (so a bare "/cmd" with only notifications still runs
+// something sensible rather than an empty prompt). Any other input — a
+// non-command, an unknown command, or a call error — leaves prompt unchanged so
+// the normal headless run proceeds. mgr may be nil (no plugins).
+//
+// Headless has no turn-injection loop, so "inject the returned prompt" degrades
+// to "use the returned prompt for this run", which the acceptance criteria
+// permit.
+func resolveHeadlessPluginCommand(prompt string, mgr *plugin.Manager, notifyOut io.Writer) string {
+	if mgr == nil || !strings.HasPrefix(strings.TrimLeft(prompt, " \t"), "/") {
+		return prompt
+	}
+	trimmed := strings.TrimLeft(prompt, " \t")[1:]
+	name := trimmed
+	args := ""
+	if i := strings.IndexAny(trimmed, " \t"); i >= 0 {
+		name = trimmed[:i]
+		args = strings.TrimSpace(trimmed[i+1:])
+	}
+	for _, pc := range mgr.Commands() {
+		if pc.Spec.Name != name {
+			continue
+		}
+		// Encode the raw arg text as a JSON string (never null), matching the
+		// host's CommandCallParams.Args contract.
+		raw, _ := json.Marshal(args)
+		res, err := pc.Plugin.CallCommand(context.Background(), name, json.RawMessage(raw))
+		if err != nil {
+			fmt.Fprintf(notifyOut, "pigo: plugin command %q failed: %v\n", name, err)
+			return prompt
+		}
+		for _, n := range res.Notifications {
+			if n.Type != "" {
+				fmt.Fprintf(notifyOut, "[%s] %s\n", n.Type, n.Message)
+			} else {
+				fmt.Fprintln(notifyOut, n.Message)
+			}
+		}
+		if res.Prompt != "" {
+			return res.Prompt
+		}
+		return args
+	}
+	return prompt
+}
+
+// ParseOutputMode maps the --output-format flag onto a HeadlessMode, erroring on
+// an unknown value.
+func ParseOutputMode(outputFmt string) (runtime.HeadlessMode, error) {
+	switch outputFmt {
+	case "text", "":
+		return runtime.PrintMode, nil
+	case "stream-json":
+		return runtime.StreamJSONMode, nil
+	default:
+		return 0, fmt.Errorf("unknown --output-format %q (want text|stream-json)", outputFmt)
+	}
+}

pigo/internal/cli/headless/headless_test.go

@@ -0,0 +1,41 @@
+package headless
+
+// Tests for the headless run driver's flag parsing. The run lifecycle (Run) is
+// exercised via session/subagent tests here and provider-backed tests in
+// internal/runtime; ParseOutputMode is pure and pinned directly.
+
+import (
+	"testing"
+
+	"github.com/smallnest/pigo/internal/runtime"
+)
+
+// TestParseOutputMode covers the three accepted spellings and one rejection,
+// pinning the flag contract the headless driver depends on.
+func TestParseOutputMode(t *testing.T) {
+	cases := []struct {
+		in      string
+		want    runtime.HeadlessMode
+		wantErr bool
+	}{
+		{"text", runtime.PrintMode, false},
+		{"", runtime.PrintMode, false},
+		{"stream-json", runtime.StreamJSONMode, false},
+		{"yaml", 0, true},
+	}
+	for _, c := range cases {
+		got, err := ParseOutputMode(c.in)
+		if c.wantErr {
+			if err == nil {
+				t.Errorf("ParseOutputMode(%q): want error, got nil", c.in)
+			}
+			continue
+		}
+		if err != nil {
+			t.Errorf("ParseOutputMode(%q): unexpected error %v", c.in, err)
+		}
+		if got != c.want {
+			t.Errorf("ParseOutputMode(%q) = %v, want %v", c.in, got, c.want)
+		}
+	}
+}

pigo/internal/cli/headless/session.go

@@ -0,0 +1,185 @@
+// Package headless drives pigo's non-interactive run paths: the print /
+// stream-json headless run, the session listing/resume helpers, and the
+// process-isolated sub-agent JSON-RPC server (--subagent-rpc).
+//
+// This file gives headless / stream-json runs the same session persistence and
+// resume the interactive REPL has (cmd/pigo/interactive.go). Before this, a
+// headless run built an in-memory AgentContext and threw it away on exit, so
+// `--output-format stream-json` emitted no session id and `--resume`/`--continue`
+// only worked in the REPL.
+//
+// Now a headless run is backed by a session file: resuming seeds the context
+// from a prior session (and re-anchors the branch leaf), a fresh run creates a
+// new session, and in both cases the run's newly produced messages are appended
+// after it completes. The session id is threaded into the run so it appears in
+// the first stream-json event (mirrors pi/Claude Code) and can be passed back via
+// --resume to continue the run.
+package headless
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+	"github.com/smallnest/pigo/internal/session"
+)
+
+// SessionStore returns the session store rooted at ~/.pigo/sessions (or under
+// PIGO_HOME when set), creating the directory on first use. It is shared by the
+// headless run path and the interactive REPL.
+func SessionStore() (*session.Store, error) {
+	dir := os.Getenv("PIGO_HOME")
+	if dir == "" {
+		home, err := os.UserHomeDir()
+		if err != nil {
+			return nil, fmt.Errorf("resolve home dir: %w", err)
+		}
+		dir = filepath.Join(home, ".pigo")
+	}
+	return session.NewStore(filepath.Join(dir, "sessions"))
+}
+
+// PrintSessions prints the stored sessions, most-recent first, to out.
+func PrintSessions(out io.Writer) error {
+	store, err := SessionStore()
+	if err != nil {
+		return err
+	}
+	headers, err := store.List()
+	if err != nil {
+		return err
+	}
+	if len(headers) == 0 {
+		fmt.Fprintln(out, "no sessions")
+		return nil
+	}
+	for _, h := range headers {
+		fmt.Fprintf(out, "%s\t%s\t%s\n", h.ID, h.UpdatedAt.Local().Format("2006-01-02 15:04"), h.Model)
+	}
+	return nil
+}
+
+// MostRecentSessionID returns the id of the most recently updated session, or
+// "" if there are none.
+func MostRecentSessionID() (string, error) {
+	store, err := SessionStore()
+	if err != nil {
+		return "", err
+	}
+	headers, err := store.List()
+	if err != nil {
+		return "", err
+	}
+	if len(headers) == 0 {
+		return "", nil
+	}
+	return headers[0].ID, nil
+}
+
+// headlessSession is the session state backing one headless run: the store, the
+// header (whose ID is the session id emitted and used for resume), and the
+// branch-tracking cursor (curLeaf/persisted) so the run's messages append as a
+// branch descending from the resumed leaf rather than flattening the tree.
+type headlessSession struct {
+	store   *session.Store
+	header  session.SessionHeader
+	curLeaf string // active leaf id to descend from; "" for a fresh session
+	// persisted is the number of agentCtx.Messages already on disk before the
+	// run; persist appends only Messages[persisted:] as a new branch.
+	persisted int
+	// model/provider are the model and provider the run actually used, refreshed
+	// onto the header before persisting so a resumed run does not write back the
+	// original session's stale values (matching the REPL, repl.go persistTurn).
+	model    string
+	provider string
+}
+
+// openHeadlessSession resolves the session backing a headless run: it resumes an
+// existing session when resumeID is set (seeding priorMsgs and re-anchoring the
+// branch leaf) or creates a fresh session header otherwise. It returns the prior
+// messages to seed into the context ahead of the new prompt, plus the session
+// state used to persist the run afterward.
+func openHeadlessSession(resumeID, model, providerName, sysPrompt string) (agentcore.MessageList, headlessSession, error) {
+	store, err := SessionStore()
+	if err != nil {
+		return nil, headlessSession{}, err
+	}
+	now := time.Now().UTC()
+
+	if resumeID != "" {
+		h, entries, err := store.LoadEntries(resumeID)
+		if err != nil {
+			return nil, headlessSession{}, err
+		}
+		msgs := make(agentcore.MessageList, len(entries))
+		for i, e := range entries {
+			msgs[i] = e.Message
+		}
+		curLeaf := ""
+		if len(entries) > 0 {
+			curLeaf = entries[len(entries)-1].ID
+		}
+		// A resumed header keeps its own SystemPrompt when present so the run is
+		// faithful to the original session.
+		if h.SystemPrompt == "" {
+			h.SystemPrompt = sysPrompt
+		}
+		return msgs, headlessSession{store: store, header: h, curLeaf: curLeaf, persisted: len(msgs), model: model, provider: providerName}, nil
+	}
+
+	header := session.SessionHeader{
+		ID:           session.NewID(now),
+		CreatedAt:    now,
+		UpdatedAt:    now,
+		Model:        model,
+		Provider:     providerName,
+		SystemPrompt: sysPrompt,
+		Cwd:          headlessCwd(),
+	}
+	return nil, headlessSession{store: store, header: header, curLeaf: "", persisted: 0, model: model, provider: providerName}, nil
+}
+
+// headlessCwd returns the absolute working directory the run executes in, used
+// to attribute the session to a project (SessionHeader.Cwd → project id) so a
+// later /dream pass can distill this session under the right project scope. An
+// unresolvable cwd yields "" (the session stays unattributed) rather than
+// aborting the run.
+func headlessCwd() string {
+	wd, err := os.Getwd()
+	if err != nil {
+		return ""
+	}
+	return wd
+}
+
+// persist appends the messages produced during the run — everything in
+// agentCtx.Messages past what was already on disk — as a branch descending from
+// the resumed leaf, matching how the REPL grows a session tree (AppendBranch).
+// It is a no-op when the run produced nothing new. Errors are returned for the
+// caller to surface; the run's output has already been emitted regardless.
+func (hs *headlessSession) persist(agentCtx *agentcore.AgentContext) error {
+	// Compaction can rebuild agentCtx.Messages to fewer entries than were on disk
+	// before the run (loop.go maybeAutoCompact replaces the slice). Clamp the
+	// cursor so the tail slice stays in bounds; when the context shrank there is
+	// nothing new to append past what compaction kept.
+	if hs.persisted > len(agentCtx.Messages) {
+		hs.persisted = len(agentCtx.Messages)
+	}
+	tail := agentCtx.Messages[hs.persisted:]
+	if len(tail) == 0 {
+		return nil
+	}
+	// Refresh the header with the model/provider the run actually used so a
+	// resumed session's metadata is not written back stale (matching the REPL).
+	hs.header.Model = hs.model
+	hs.header.Provider = hs.provider
+	hs.header.UpdatedAt = time.Now().UTC()
+	if _, err := hs.store.AppendBranch(hs.header, hs.curLeaf, tail); err != nil {
+		return err
+	}
+	hs.persisted = len(agentCtx.Messages)
+	return nil
+}

pigo/internal/cli/headless/session_test.go

@@ -0,0 +1,156 @@
+package headless
+
+// Tests for headless session persistence and resume (session id in stream-json
+// + --resume for headless runs). openHeadlessSession/persist are exercised
+// directly against an isolated PIGO_HOME so a headless run's session round-trips
+// without spawning a provider.
+
+import (
+	"testing"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+)
+
+func textUser(s string) agentcore.UserMessage {
+	return agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent(s)}}
+}
+
+func textAssistant(s string) agentcore.AssistantMessage {
+	return agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent(s)}}
+}
+
+// TestOpenHeadlessSessionFresh verifies a fresh headless session gets a new id
+// and empty prior messages, and that persist writes the run's messages so they
+// can be resumed.
+func TestOpenHeadlessSessionFresh(t *testing.T) {
+	t.Setenv("PIGO_HOME", t.TempDir())
+
+	prior, hs, err := openHeadlessSession("", "faux-model", "faux", "sys prompt")
+	if err != nil {
+		t.Fatalf("openHeadlessSession fresh: %v", err)
+	}
+	if len(prior) != 0 {
+		t.Errorf("fresh session must have no prior messages, got %d", len(prior))
+	}
+	if hs.header.ID == "" {
+		t.Fatal("fresh session must have a non-empty id")
+	}
+	if hs.header.SystemPrompt != "sys prompt" {
+		t.Errorf("header SystemPrompt = %q, want the passed prompt", hs.header.SystemPrompt)
+	}
+
+	// Simulate a completed run: prompt + assistant reply.
+	agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("1+1=?"), textAssistant("2")}}
+	if err := hs.persist(agentCtx); err != nil {
+		t.Fatalf("persist: %v", err)
+	}
+
+	// The session must now be loadable with both messages.
+	_, msgs, err := hs.store.Load(hs.header.ID)
+	if err != nil {
+		t.Fatalf("Load persisted session: %v", err)
+	}
+	if len(msgs) != 2 {
+		t.Fatalf("persisted session has %d messages, want 2", len(msgs))
+	}
+}
+
+// TestOpenHeadlessSessionResume verifies that resuming seeds the prior messages
+// and that a subsequent run appends only the new tail as a branch, so the
+// session grows rather than being rewritten.
+func TestOpenHeadlessSessionResume(t *testing.T) {
+	t.Setenv("PIGO_HOME", t.TempDir())
+
+	// First run: create and persist a session.
+	_, hs1, err := openHeadlessSession("", "faux-model", "faux", "sys")
+	if err != nil {
+		t.Fatalf("first openHeadlessSession: %v", err)
+	}
+	ctx1 := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("first"), textAssistant("reply1")}}
+	if err := hs1.persist(ctx1); err != nil {
+		t.Fatalf("first persist: %v", err)
+	}
+	sessID := hs1.header.ID
+
+	// Second run: resume the session id.
+	prior, hs2, err := openHeadlessSession(sessID, "faux-model", "faux", "sys")
+	if err != nil {
+		t.Fatalf("resume openHeadlessSession: %v", err)
+	}
+	if len(prior) != 2 {
+		t.Fatalf("resume must seed %d prior messages, got %d", 2, len(prior))
+	}
+	if hs2.header.ID != sessID {
+		t.Errorf("resumed session id = %q, want %q", hs2.header.ID, sessID)
+	}
+	if hs2.persisted != 2 {
+		t.Errorf("resumed persisted cursor = %d, want 2", hs2.persisted)
+	}
+
+	// A second turn appends its new tail.
+	ctx2 := &agentcore.AgentContext{Messages: append(prior, textUser("second"), textAssistant("reply2"))}
+	if err := hs2.persist(ctx2); err != nil {
+		t.Fatalf("second persist: %v", err)
+	}
+	_, msgs, err := hs2.store.Load(sessID)
+	if err != nil {
+		t.Fatalf("Load after second turn: %v", err)
+	}
+	if len(msgs) != 4 {
+		t.Fatalf("session after two turns has %d messages, want 4", len(msgs))
+	}
+}
+
+// TestHeadlessPersistNoop verifies persist is a no-op (no error, no growth) when
+// the run produced nothing new past what was already persisted.
+func TestHeadlessPersistNoop(t *testing.T) {
+	t.Setenv("PIGO_HOME", t.TempDir())
+	_, hs, err := openHeadlessSession("", "m", "p", "s")
+	if err != nil {
+		t.Fatalf("openHeadlessSession: %v", err)
+	}
+	ctx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("x"), textAssistant("y")}}
+	if err := hs.persist(ctx); err != nil {
+		t.Fatalf("first persist: %v", err)
+	}
+	// Persisting again with no new messages must not error and must not duplicate.
+	if err := hs.persist(ctx); err != nil {
+		t.Fatalf("noop persist: %v", err)
+	}
+	_, msgs, err := hs.store.Load(hs.header.ID)
+	if err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+	if len(msgs) != 2 {
+		t.Fatalf("noop persist changed message count to %d, want 2", len(msgs))
+	}
+}
+
+// TestHeadlessPersistCompactionShrink verifies persist tolerates the context
+// being rebuilt to fewer messages than were on disk before the run (mid-run
+// compaction replaces agentCtx.Messages). The persisted cursor is clamped so
+// the tail slice stays in bounds rather than panicking.
+func TestHeadlessPersistCompactionShrink(t *testing.T) {
+	t.Setenv("PIGO_HOME", t.TempDir())
+	_, hs, err := openHeadlessSession("", "m", "p", "s")
+	if err != nil {
+		t.Fatalf("openHeadlessSession: %v", err)
+	}
+	// Persist four messages, advancing the cursor to 4.
+	ctx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("a"), textAssistant("b"), textUser("c"), textAssistant("d")}}
+	if err := hs.persist(ctx); err != nil {
+		t.Fatalf("first persist: %v", err)
+	}
+	if hs.persisted != 4 {
+		t.Fatalf("cursor = %d, want 4", hs.persisted)
+	}
+	// Simulate compaction: the context is rebuilt to fewer messages than the
+	// cursor. persist must not panic on the out-of-range slice.
+	ctx.Messages = agentcore.MessageList{textAssistant("summary"), textUser("e")}
+	if err := hs.persist(ctx); err != nil {
+		t.Fatalf("persist after compaction shrink: %v", err)
+	}
+	if hs.persisted != 2 {
+		t.Errorf("cursor after clamp = %d, want 2", hs.persisted)
+	}
+}

pigo/internal/cli/headless/subagent_rpc.go

@@ -0,0 +1,161 @@
+// This file implements the subprocess side of process-isolated sub-agents
+// (US-019, #135). Invoked as `pigo --subagent-rpc`, pigo speaks JSON-RPC 2.0
+// over stdio: for each "subagent/run" request on stdin it runs a child agent
+// loop and writes the result (or an error) to stdout, exiting when stdin
+// closes. The parent (SubAgentTool in process mode, internal/runtime) drives it
+// via internal/jsonrpc.
+//
+// It reuses internal/jsonrpc's message types (Request/Response/ID/Version) for
+// (de)serialization so the wire format matches the client exactly; the agent
+// execution itself is runtime.RunSubAgentOnce.
+package headless
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+	"github.com/smallnest/pigo/internal/agenttool"
+	"github.com/smallnest/pigo/internal/cli/run"
+	"github.com/smallnest/pigo/internal/jsonrpc"
+	"github.com/smallnest/pigo/internal/provider"
+	"github.com/smallnest/pigo/internal/runtime"
+)
+
+// RunSubAgentRPC is the `pigo --subagent-rpc` entry point. It reads
+// newline-delimited JSON-RPC requests from in, runs each sub-agent request, and
+// writes one response per request to out. It returns 0 (success) when stdin
+// closes; a per-request failure is an RPC error response, not a non-zero exit,
+// so the parent can distinguish "the child answered with an error" from "the
+// child crashed" (the latter is detected by the parent's transport when stdout
+// closes without a response).
+func RunSubAgentRPC(ctx context.Context, in io.Reader, out, errOut io.Writer) int {
+	scanner := bufio.NewScanner(in)
+	// A sub-agent prompt can be large; match the jsonrpc client's line cap.
+	scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
+	enc := json.NewEncoder(out)
+	for scanner.Scan() {
+		line := bytes.TrimSpace(scanner.Bytes())
+		if len(line) == 0 {
+			continue
+		}
+		var req jsonrpc.Request
+		if err := json.Unmarshal(line, &req); err != nil {
+			// A parse error carries no id, so the response id is null. The
+			// jsonrpc client drops responses with a null id (it cannot correlate
+			// them), so this is only observable on the child's stderr; in
+			// practice the parent always sends well-formed requests.
+			writeSubAgentError(enc, nil, -32700, "parse error: "+err.Error())
+			continue
+		}
+		handleSubAgentRequest(ctx, enc, &req)
+	}
+	// A scanner error (e.g. a request line exceeding the 16 MiB cap) ends the
+	// stream abnormally: surface it on stderr and exit non-zero so the parent's
+	// transport sees a diagnostic rather than a silent clean exit.
+	if err := scanner.Err(); err != nil {
+		fmt.Fprintf(errOut, "pigo: subagent-rpc stdin: %v\n", err)
+		return 1
+	}
+	return 0
+}
+
+// handleSubAgentRequest dispatches one JSON-RPC request to the sub-agent runner
+// and writes the response. Unknown methods, bad params, provider-resolution
+// failures, and failed child runs are all RPC errors so the parent surfaces them
+// as tool errors; only a successful run yields a result with the child's text.
+func handleSubAgentRequest(ctx context.Context, enc *json.Encoder, req *jsonrpc.Request) {
+	if req.Method != runtime.SubAgentRPCMethod {
+		writeSubAgentError(enc, req.ID, -32601, "method not found: "+req.Method)
+		return
+	}
+	var params runtime.SubAgentRunParams
+	if len(req.Params) > 0 {
+		if err := json.Unmarshal(req.Params, &params); err != nil {
+			writeSubAgentError(enc, req.ID, -32602, "invalid params: "+err.Error())
+			return
+		}
+	}
+	if params.Prompt == "" || params.Model == "" {
+		writeSubAgentError(enc, req.ID, -32602, "invalid params: prompt and model are required")
+		return
+	}
+	// Resolve the provider the same way the CLI does, so the subprocess targets
+	// the same gateway the parent's NewRunConfig encoded. Credentials come from
+	// the inherited environment (the parent's env vars).
+	prov, providerName, err := provider.ResolveProvider(params.Model, params.BaseURL, params.Protocol, "", os.Getenv)
+	if err != nil {
+		writeSubAgentError(enc, req.ID, -32603, "resolve provider: "+err.Error())
+		return
+	}
+	cwd, _ := os.Getwd()
+	tools := filterBuiltinTools(run.BuiltinTools(cwd, false), params.Tools)
+	reg := run.ToolRegistry(tools)
+	creds := provider.NewCredentialStore(nil) // env-resolved
+	runCfg := runtime.RunConfig{
+		LoopConfig: runtime.LoopConfig{
+			Model:     params.Model,
+			Provider:  providerName,
+			Stream:    provider.StreamFnFromProvider(prov),
+			GetAPIKey: creds.GetAPIKey,
+		},
+		Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}},
+	}
+	// Wire hooks uniformly with every other driver (#425): the child sub-agent runs
+	// its own PreToolUse/PostToolUse (and Stop) hooks from the trust-gated hook set
+	// rooted at its working directory. It has no backing session, so SessionID is
+	// empty (omitted from HookInput). A malformed hook layer disables hooks with a
+	// warning rather than failing the child run.
+	if set, herr := run.ResolveHookSet(cwd, run.Trusted(cwd)); herr != nil {
+		fmt.Fprintf(os.Stderr, "pigo: hooks disabled: %v\n", herr)
+	} else {
+		run.InstallHooks(&runCfg, set, run.HookDeps{ProjectDir: cwd, WarnLog: os.Stderr})
+	}
+	text, err := runtime.RunSubAgentOnce(ctx, params.SystemPrompt, params.Prompt, tools, runCfg)
+	if err != nil {
+		// A failed child run is an RPC error so the parent's defaultProcessCall
+		// returns a Go error and executeProcess marks the tool result IsError,
+		// matching goroutine mode's "failed run -> tool error" behavior.
+		writeSubAgentError(enc, req.ID, -32000, err.Error())
+		return
+	}
+	result, _ := json.Marshal(runtime.SubAgentRunResult{Text: text})
+	_ = enc.Encode(jsonrpc.Response{JSONRPC: jsonrpc.Version, ID: req.ID, Result: result})
+}
+
+// writeSubAgentError writes a JSON-RPC error response with the given id (which
+// may be nil for a parse error on an unidentifiable request) and code/message.
+func writeSubAgentError(enc *json.Encoder, id *jsonrpc.ID, code int, msg string) {
+	_ = enc.Encode(jsonrpc.Response{
+		JSONRPC: jsonrpc.Version,
+		ID:      id,
+		Error:   &jsonrpc.Error{Code: code, Message: msg},
+	})
+}
+
+// filterBuiltinTools returns the subset of tools whose Name is in names. An
+// empty names list keeps all tools. It lets a process-isolated sub-agent
+// restrict the child to a subset (e.g. a read-only researcher), matching
+// goroutine mode's Tools filtering, without serializing in-process tool objects
+// across the process boundary.
+func filterBuiltinTools(tools []agentcore.AgentTool, names []string) []agentcore.AgentTool {
+	if len(names) == 0 {
+		return tools
+	}
+	want := make(map[string]bool, len(names))
+	for _, n := range names {
+		want[n] = true
+	}
+	var out []agentcore.AgentTool
+	for _, t := range tools {
+		if want[t.Name()] {
+			out = append(out, t)
+		}
+	}
+	return out
+}

pigo/internal/cli/headless/subagent_rpc_test.go

@@ -0,0 +1,120 @@
+package headless
+
+// Tests for the sub-agent RPC subprocess mode (US-019, #135): the pure
+// filterBuiltinTools helper and the RunSubAgentRPC validation branches
+// (method-not-found, invalid params, parse error) that return RPC errors before
+// any provider is resolved. The happy-path transport is covered in
+// internal/runtime via a compiled helper binary; RunSubAgentOnce (the agent
+// core) is covered there with a faux provider.
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+	"github.com/smallnest/pigo/internal/cli/run"
+	"github.com/smallnest/pigo/internal/jsonrpc"
+	"github.com/smallnest/pigo/internal/runtime"
+)
+
+// TestFilterBuiltinTools verifies the subprocess tool filter: an empty name
+// list keeps all builtins, a subset keeps only the named tools, and unknown
+// names are silently ignored.
+func TestFilterBuiltinTools(t *testing.T) {
+	all := run.BuiltinTools(t.TempDir(), false)
+	if len(all) == 0 {
+		t.Fatal("BuiltinTools returned no tools")
+	}
+	namesOf := func(ts []agentcore.AgentTool) []string {
+		out := make([]string, len(ts))
+		for i, tl := range ts {
+			out[i] = tl.Name()
+		}
+		return out
+	}
+
+	if got := filterBuiltinTools(all, nil); len(got) != len(all) {
+		t.Errorf("nil names kept %d, want all %d", len(got), len(all))
+	}
+	if got := filterBuiltinTools(all, []string{}); len(got) != len(all) {
+		t.Errorf("empty names kept %d, want all %d", len(got), len(all))
+	}
+
+	got := filterBuiltinTools(all, []string{"read", "grep"})
+	if len(got) != 2 {
+		t.Fatalf("subset kept %d, want 2: %v", len(got), namesOf(got))
+	}
+	gotNames := namesOf(got)
+	if gotNames[0] != "read" || gotNames[1] != "grep" {
+		t.Errorf("subset names = %v, want [read grep]", gotNames)
+	}
+
+	// Unknown names are ignored, known ones kept.
+	got = filterBuiltinTools(all, []string{"read", "does-not-exist"})
+	if len(got) != 1 || got[0].Name() != "read" {
+		t.Errorf("unknown-name filter kept %v, want [read]", namesOf(got))
+	}
+}
+
+// TestRunSubAgentRPCValidation verifies the subprocess returns JSON-RPC errors
+// for malformed requests without reaching provider resolution: a parse error,
+// an unknown method, and missing prompt/model params. These branches are the
+// server's contract for bad input and need not involve a provider.
+func TestRunSubAgentRPCValidation(t *testing.T) {
+	cases := []struct {
+		name     string
+		line     string
+		wantCode int
+		wantMsg  string
+	}{
+		{"parse error", "not-json", -32700, "parse error"},
+		{"method not found", `{"jsonrpc":"2.0","id":1,"method":"other","params":{}}`, -32601, "method not found"},
+		{"missing prompt", `{"jsonrpc":"2.0","id":2,"method":"` + runtime.SubAgentRPCMethod + `","params":{"model":"x"}}`, -32602, "prompt and model are required"},
+		{"missing model", `{"jsonrpc":"2.0","id":3,"method":"` + runtime.SubAgentRPCMethod + `","params":{"prompt":"x"}}`, -32602, "prompt and model are required"},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			var out, errOut bytes.Buffer
+			code := RunSubAgentRPC(context.Background(), strings.NewReader(c.line+"\n"), &out, &errOut)
+			if code != 0 {
+				t.Errorf("exit code = %d, want 0 (validation errors are RPC responses, not non-zero exits)", code)
+			}
+			line, _ := out.ReadString('\n')
+			if line == "" {
+				t.Fatal("no response written")
+			}
+			var resp jsonrpc.Response
+			if err := json.Unmarshal([]byte(line), &resp); err != nil {
+				t.Fatalf("unmarshal response %q: %v", line, err)
+			}
+			if resp.Error == nil {
+				t.Fatalf("response has no error: %s", line)
+			}
+			if resp.Error.Code != c.wantCode {
+				t.Errorf("error code = %d, want %d (msg=%q)", resp.Error.Code, c.wantCode, resp.Error.Message)
+			}
+			if !strings.Contains(resp.Error.Message, c.wantMsg) {
+				t.Errorf("error msg = %q, want it to contain %q", resp.Error.Message, c.wantMsg)
+			}
+		})
+	}
+}
+
+// TestRunSubAgentRPCScannerError verifies a stdin read error (a line exceeding
+// the scanner cap) is surfaced on stderr and yields a non-zero exit, rather
+// than a silent clean exit.
+func TestRunSubAgentRPCScannerError(t *testing.T) {
+	// A line longer than the 16 MiB scanner cap triggers a scanner error.
+	huge := strings.Repeat("a", 17*1024*1024)
+	var out, errOut bytes.Buffer
+	code := RunSubAgentRPC(context.Background(), strings.NewReader(huge), &out, &errOut)
+	if code == 0 {
+		t.Error("exit code = 0 on scanner error, want non-zero")
+	}
+	if !strings.Contains(errOut.String(), "subagent-rpc stdin") {
+		t.Errorf("stderr = %q, want it to mention 'subagent-rpc stdin'", errOut.String())
+	}
+}