Vectrum
Sign In
261 lines · 8.5 KB · Go
package main

// The command line another program drives Codey through.
//
// Three invocations here are not for a person. A control plane that runs
// coding agents — T3 Code is the one this was built against — spawns the
// binary three ways: once to ask its version, once to ask what models it has
// and whether it is logged in, and once to run the agent itself over a pipe.
// Only the third does any work; the first two exist so the provider card can
// say "installed" and "authenticated" before anyone starts a thread.
//
// The shapes are not ours to choose. They are what the client already sends to
// the binary whose slot Codey is occupying, which is why the flags arrive
// around the subcommand rather than before it and why the two probes print
// prose a regular expression reads.

import (
	"context"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"time"

	"codey/internal/acp"
	"codey/internal/agent"
	"codey/internal/llm"
	"codey/internal/metrics"
	"codey/internal/serve"
	"codey/internal/tools"
)

// invKind is which of the four things this command line is.
type invKind int

const (
	invNormal invKind = iota
	invVersion
	invModels
	invAgent
)

// invocation is the command line, read.
type invocation struct {
	Kind invKind
	// Permission is the approval floor an agent run was launched with.
	Permission serve.PermissionMode
	// Rest is what is left for ordinary flag parsing once the words that
	// belong to the subcommand are taken out.
	Rest []string
}

// scanArgs reads the command line before the flag package sees it.
//
// It has to, for two reasons. `--version` is not one of our flags, and flag
// parsing would exit(2) on it rather than answer. And the agent invocation
// splits its flags around the subcommand — `--permission-mode default agent
// stdio`, but `agent --always-approve stdio` — which no flag package parses,
// because it is not one grammar but a client's habit.
//
// Everything unrecognised is handed on rather than rejected. The launch
// arguments are a free-text field in somebody's settings panel, and the flags
// Codey does understand should keep working when they arrive through it.
func scanArgs(args []string) invocation {
	// The probes are exact: the client runs the binary with this one word and
	// nothing else. Requiring that keeps `-C models` from being mistaken for
	// a request to list them.
	if len(args) == 1 {
		switch args[0] {
		case "--version", "-version":
			return invocation{Kind: invVersion}
		case "models":
			return invocation{Kind: invModels}
		}
	}

	var (
		sawAgent, sawStdio, always bool
		mode                       string
		rest                       []string
	)
	for i := 0; i < len(args); i++ {
		a := args[i]
		switch {
		case a == "agent":
			sawAgent = true
		case a == "stdio":
			sawStdio = true
		case a == "--always-approve", a == "-always-approve":
			always = true
		case a == "--permission-mode", a == "-permission-mode":
			if i+1 < len(args) {
				mode = args[i+1]
				i++
			}
		case strings.HasPrefix(a, "--permission-mode="):
			mode = strings.TrimPrefix(a, "--permission-mode=")
		case strings.HasPrefix(a, "-permission-mode="):
			mode = strings.TrimPrefix(a, "-permission-mode=")
		default:
			rest = append(rest, a)
		}
	}

	if !sawAgent || !sawStdio {
		// Not an agent run after all, so nothing was ours to consume: hand the
		// line back exactly as it came.
		return invocation{Kind: invNormal, Rest: args}
	}
	return invocation{Kind: invAgent, Permission: permissionMode(mode, always), Rest: rest}
}

// permissionMode maps the client's vocabulary onto ours. An unrecognised mode
// is the strictest one: a name we do not know is not a reason to stop asking.
func permissionMode(mode string, always bool) serve.PermissionMode {
	if always {
		return serve.ModeFullAccess
	}
	switch mode {
	case "acceptEdits":
		return serve.ModeAcceptEdits
	case "auto":
		return serve.ModeAuto
	case "bypassPermissions", "full-access":
		return serve.ModeFullAccess
	default:
		return serve.ModeAsk
	}
}

// printVersion answers the version probe.
//
// The line has to contain something matching \bv?(\d+\.\d+\.\d+)\b, because
// that is the regular expression on the other end. Anything else reads as an
// installed binary that cannot run.
func printVersion(w io.Writer) {
	fmt.Fprintf(w, "codey %s\n", version)
}

// modelsTimeout bounds the listing. The probe that runs this has its own,
// shorter patience, and a listing that has not arrived by then is better
// reported as the one model we know about than not reported at all.
const modelsTimeout = 8 * time.Second

// printModels answers the model-and-login probe.
//
// The format is prose because the reader is a regular expression: a line
// saying whether we are logged in, then one bullet per model with the default
// marked. It exits zero even when it has nothing good to say — a non-zero exit
// makes the client discard the whole listing, including the part that was
// true.
func printModels(w io.Writer, client *llm.Client) {
	if !client.HasCredentials() {
		fmt.Fprintln(w, "You are not logged in. Set ANTHROPIC_API_KEY, or log in with Claude Code.")
		return
	}
	fmt.Fprintf(w, "You are logged in to %s.\n", client.BaseURL)

	ctx, cancel := context.WithTimeout(context.Background(), modelsTimeout)
	defer cancel()

	models, err := client.ListModels(ctx)
	if err != nil || len(models) == 0 {
		// An endpoint that does not answer still serves the model we were
		// configured with, and offering that beats offering nothing.
		fmt.Fprintf(w, "* %s (default)\n", client.Model())
		return
	}

	for _, m := range sortedModels(models) {
		if m.ID == client.Model() {
			fmt.Fprintf(w, "* %s (default)\n", m.ID)
			continue
		}
		fmt.Fprintf(w, "- %s\n", m.ID)
	}
}

// sortedModels puts the newest first, which is the order a picker should show
// and not the order every endpoint answers in.
func sortedModels(models []llm.Model) []llm.Model {
	out := append([]llm.Model(nil), models...)
	sort.SliceStable(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) })
	return out
}

// acpDeps is what an agent run borrows from the ordinary startup path. Nothing
// here is built for this mode: it is the same client, the same event log, the
// same options and the same tool factory the TUI gets, which is the point —
// an agent driven through a pipe should be able to do exactly what an agent
// driven through a keyboard can.
type acpDeps struct {
	Permission serve.PermissionMode
	Client     *llm.Client
	// ClientOpts rebuilds the client for another model. The model lives on the
	// client rather than on the agent, so a session that selects one needs a
	// client of its own.
	ClientOpts []llm.Option
	Log        *metrics.Logger
	Base       agent.Options
	// Tools is main's factory: a root in, a tool-set builder out.
	Tools func(root string) func(*agent.State, *metrics.Logger) *tools.Registry
}

// serveACP runs the agent over stdin and stdout until the client hangs up.
//
// Nothing else in this process may write to stdout from here on. The stream is
// the protocol, and one stray line desynchronises a parser that has no way to
// resynchronise.
func serveACP(ctx context.Context, deps acpDeps) error {
	models, current := offeredModels(ctx, deps.Client)

	srv := serve.New(serve.Options{
		Version:      version,
		Permission:   deps.Permission,
		Models:       models,
		DefaultModel: current,

		NewAgent: func(root, model string, state *agent.State) (*agent.Agent, error) {
			client := deps.Client
			if model != "" && model != client.Model() {
				client = llm.New(append(deps.ClientOpts, llm.WithModel(model))...)
			}
			reg := deps.Tools(root)(state, deps.Log)
			return agent.New(client, reg, state, deps.Log, deps.Base), nil
		},
	})

	deps.Log.Log(metrics.Event{Type: "acp_start", Extra: map[string]any{
		"permission": string(deps.Permission), "model": deps.Client.Model(),
	}})

	return srv.Run(ctx, os.Stdin, os.Stdout)
}

// offeredModels is the picker a session is given. An endpoint that will not
// list them leaves the picker empty, and the client then shows whatever it
// learned from the `models` probe instead.
func offeredModels(ctx context.Context, client *llm.Client) ([]acp.ModelInfo, string) {
	listCtx, cancel := context.WithTimeout(ctx, modelsTimeout)
	defer cancel()

	found, err := client.ListModels(listCtx)
	if err != nil || len(found) == 0 {
		return nil, client.Model()
	}

	out := make([]acp.ModelInfo, 0, len(found))
	for _, m := range sortedModels(found) {
		out = append(out, acp.ModelInfo{ModelID: m.ID, Name: displayName(m)})
	}
	return out, client.Model()
}

func displayName(m llm.Model) string {
	if m.DisplayName != "" {
		return m.DisplayName
	}
	return m.ID
}