Vectrum
Sign In
301 lines · 9.7 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/subagent"
	"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
	// NoAgents leaves the subagent tools out, so the agent works alone.
	NoAgents bool
	// 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 {
	found := listModels(ctx, deps.Client)

	// What a subagent runs under is its parent's settings without the model
	// listing: a child never chooses a model for a child of its own, and a
	// paragraph about a choice it will not make is one it pays for on every
	// request. The parent does get the listing, because it is the one that
	// chooses.
	subBase := deps.Base
	base := deps.Base
	base.Models = found

	srv := serve.New(serve.Options{
		Version:      version,
		Permission:   deps.Permission,
		Models:       offeredModels(found),
		DefaultModel: deps.Client.Model(),

		NewAgent: func(req serve.AgentRequest) (serve.AgentSet, error) {
			client := deps.Client
			if req.Model != "" && req.Model != client.Model() {
				client = llm.New(append(deps.ClientOpts, llm.WithModel(req.Model))...)
			}

			// One factory for the session and for every child of it: an agent
			// that could reach further than its parent would be a way around
			// the parent's limits.
			newTools := deps.Tools(req.Root)
			reg := newTools(req.State, deps.Log)

			set := serve.AgentSet{}
			if !deps.NoAgents {
				sup := subagent.New(ctx, subagent.Options{
					Client: client,
					Log:    deps.Log,
					Policy: req.Approve,
					Base:   subBase,
					Tools:  newTools,
				})
				reg.Register(sup.Tools()...)
				set.Agents = sup
			}

			set.Agent = agent.New(client, reg, req.State, deps.Log, base)
			return set, 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)
}

// listModels asks the endpoint what it serves. An endpoint that will not say
// leaves the list empty, which is the honest answer: no list beats a made-up
// one, and both the picker and the agent's own model choices degrade to the
// model this process was configured with.
func listModels(ctx context.Context, client *llm.Client) []llm.Model {
	listCtx, cancel := context.WithTimeout(ctx, modelsTimeout)
	defer cancel()

	found, err := client.ListModels(listCtx)
	if err != nil {
		return nil
	}
	return sortedModels(found)
}

// offeredModels is that list as the client's model picker wants it.
func offeredModels(found []llm.Model) []acp.ModelInfo {
	if len(found) == 0 {
		return nil
	}
	out := make([]acp.ModelInfo, 0, len(found))
	for _, m := range found {
		out = append(out, acp.ModelInfo{ModelID: m.ID, Name: displayName(m)})
	}
	return out
}

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