Vectrum
Sign In
217 lines · 6.9 KB · Go
// Command harness is the CLI entry point: flag parsing, wiring, and the REPL.
package main

import (
	"bufio"
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"

	"harness/internal/agent"
	"harness/internal/llm"
	"harness/internal/metrics"
	"harness/internal/tools"
	"harness/internal/ui"
)

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, "harness:", err)
		os.Exit(1)
	}
}

func run() error {
	workdir := flag.String("C", ".", "working root the agent operates in")
	model := flag.String("model", "", "model id (default "+llm.DefaultModel+")")
	baseURL := flag.String("base-url", "", "API endpoint, for a proxy, gateway, or local server speaking the Messages API (env ANTHROPIC_BASE_URL, default "+llm.DefaultBaseURL+")")
	apiKey := flag.String("api-key", "", "credential sent as x-api-key (env ANTHROPIC_API_KEY). Visible to other processes; prefer the env var")
	authToken := flag.String("auth-token", "", "credential sent as Authorization: Bearer, for OAuth tokens and most gateways (env ANTHROPIC_AUTH_TOKEN)")
	effort := flag.String("effort", "", "output_config.effort: low|medium|high|xhigh|max (default: the API's, high)")
	maxTurns := flag.Int("max-turns", agent.DefaultMaxTurns, "turn cap per prompt")
	contextWindow := flag.Int("context-window", agent.DefaultContextWindow, "model context window, used to decide when to compact")
	compactAt := flag.Float64("compact-at", agent.DefaultCompactAt, "fraction of the window at which to compact")
	noCache := flag.Bool("no-cache", false, "disable prompt cache breakpoints")
	noStream := flag.Bool("no-stream", false, "wait for each complete response instead of streaming")
	logPath := flag.String("metrics", metrics.DefaultPath, "JSONL event log path, relative to the working root")
	prompt := flag.String("p", "", "run one prompt and exit, instead of starting the REPL")
	resume := flag.String("resume", "", "resume a saved session by id, or \"latest\"")
	verbose := flag.Bool("v", false, "print full tool output")
	flag.Parse()

	root, err := filepath.Abs(*workdir)
	if err != nil {
		return err
	}

	client := llm.New(
		llm.WithModel(*model),
		llm.WithBaseURL(*baseURL),
		llm.WithAPIKey(*apiKey),
		llm.WithAuthToken(*authToken),
	)
	if !client.HasCredentials() {
		return errNoCredentials
	}

	session := time.Now().UTC().Format("20060102T150405Z")
	log, err := metrics.Open(filepath.Join(root, *logPath), session)
	if err != nil {
		return err
	}
	defer log.Close()
	log.Log(metrics.Event{Type: "session_start", Extra: map[string]any{
		"root": root, "model": client.Model, "base_url": client.BaseURL,
	}})
	if !client.IsDefaultEndpoint() {
		// Worth saying out loud: the prompt and every file the agent reads go
		// to this endpoint, not to Anthropic.
		fmt.Fprintf(os.Stderr, "harness: sending requests to %s\n", client.BaseURL)
	}

	state := agent.NewState()
	reg := tools.NewRegistry(log).Register(
		tools.NewRead(root, state),
		tools.NewGlob(root),
		tools.NewGrep(root),
		tools.NewEdit(root, state),
		tools.NewWrite(root, state),
		tools.NewBash(root),
	)

	render := &ui.Renderer{W: os.Stdout, Verbose: *verbose}
	a := agent.New(client, reg, state, log, agent.Options{
		MaxTurns:      *maxTurns,
		Effort:        *effort,
		ContextWindow: *contextWindow,
		CompactAt:     *compactAt,
		NoCache:       *noCache,
		Stream:        !*noStream,
	})
	a.Observe = render.Handle

	sessionDir := filepath.Join(root, agent.SessionDir)
	if *resume != "" {
		prev, err := agent.ResolveSession(sessionDir, *resume)
		if err != nil {
			return err
		}
		a.Restore(prev)
		session = prev.ID
		fmt.Printf("resumed session %s (%d messages, %d turns)\n", prev.ID, len(prev.Messages), prev.Turns)
	}

	ctx, stopSignals := interruptible(context.Background())
	defer stopSignals()

	if *prompt != "" {
		return once(ctx, a, render, sessionDir, session, root, client, *prompt)
	}

	fmt.Printf("harness · %s · %s\ntools: %s\ntype a prompt, ctrl-c to interrupt a turn, ctrl-d to exit\n\n",
		client.Model, root, strings.Join(toolNames(reg), ", "))

	in := bufio.NewScanner(os.Stdin)
	in.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for {
		fmt.Print("> ")
		if !in.Scan() {
			fmt.Println()
			return in.Err()
		}
		line := strings.TrimSpace(in.Text())
		switch {
		case line == "":
			continue
		case line == "/exit", line == "/quit":
			return nil
		}
		if err := once(ctx, a, render, sessionDir, session, root, client, line); err != nil {
			render.Error(err)
		}
		fmt.Println()
	}
}

// once runs a single prompt, prints the session totals, and saves the session.
func once(ctx context.Context, a *agent.Agent, render *ui.Renderer, sessionDir, id, root string, client *llm.Client, prompt string) error {
	turnCtx, endTurn := beginTurn(ctx)
	_, runErr := a.Run(turnCtx, prompt)
	endTurn()

	render.Flush()
	usage, turns := a.State().Usage()
	render.Usage(usage, turns, client.Model)

	if err := a.Save(sessionDir, id, root, client.Model); err != nil {
		fmt.Fprintln(os.Stderr, "harness: saving session:", err)
	}
	if errors.Is(runErr, context.Canceled) {
		// An interrupted turn is not an error: the point is to leave the
		// session intact so the next message can correct course.
		fmt.Println("interrupted — the session is intact; say what to do differently")
		return nil
	}
	return runErr
}

// interrupt plumbing. The first ctrl-c cancels the turn in progress; a second
// within two seconds, or one with nothing running, exits.
var (
	turnCancel atomic.Pointer[context.CancelFunc]
	lastSignal atomic.Int64
)

func interruptible(parent context.Context) (context.Context, func()) {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	go func() {
		for range ch {
			now := time.Now()
			prev := lastSignal.Swap(now.UnixNano())
			cancel := turnCancel.Load()

			if cancel == nil || (prev != 0 && now.Sub(time.Unix(0, prev)) < 2*time.Second) {
				fmt.Fprintln(os.Stderr, "\nharness: exiting")
				os.Exit(130)
			}
			fmt.Fprintln(os.Stderr, "\nharness: interrupting this turn (ctrl-c again to exit)")
			(*cancel)()
		}
	}()
	return parent, func() { signal.Stop(ch) }
}

// beginTurn makes the running turn cancellable by the signal handler.
func beginTurn(parent context.Context) (context.Context, func()) {
	ctx, cancel := context.WithCancel(parent)
	turnCancel.Store(&cancel)
	return ctx, func() {
		turnCancel.Store(nil)
		lastSignal.Store(0)
		cancel()
	}
}

// errNoCredentials spells out every way to supply one, because "unauthorized"
// against a gateway that ignores the header is a confusing failure.
var errNoCredentials = errors.New(`no credentials. Supply one of:
    export ANTHROPIC_API_KEY=sk-ant-...      (sent as x-api-key)
    export ANTHROPIC_AUTH_TOKEN=...          (sent as Authorization: Bearer)
    -api-key / -auth-token                   (flags override the environment)
To use a non-Anthropic endpoint, set ANTHROPIC_BASE_URL or -base-url as well.`)

func toolNames(reg *tools.Registry) []string {
	var names []string
	for _, t := range reg.List() {
		names = append(names, t.Name())
	}
	return names
}