Vectrum
Sign In
307 lines · 9.7 KB · Go
package approval

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"sync"

	"github.com/johannSo/codey/internal/hooks"
)

// Decision is the outcome of an approval prompt.
type Decision int

const (
	Deny Decision = iota
	Once
	AlwaysThisCommand
	AlwaysThisTool
)

// Request is one tool call awaiting a decision.
type Request struct {
	Tool string
	// Summary is a one-line description of what will happen.
	Summary string
	// Detail is the argument the user most needs to see: the command, the
	// path, or the replacement text.
	Detail string
	// Rememberable is false when "always allow this exact command" would be a
	// lie — an opaque shell command does not determine what it runs.
	Rememberable bool
	// Agent is the subagent that wants this, empty for the session's own
	// agent. Several agents share one prompt box, so a request that does not
	// say who is asking cannot be decided on.
	Agent string
}

// Prompt asks a human. Returning an error aborts the turn.
type Prompt func(context.Context, Request) (Decision, error)

// readOnly tools cannot change anything, so they never prompt.
//
// web_search and web_fetch are deliberately absent. They change nothing on
// disk, but they send words the model chose — out of everything the model has
// read — to a third party, and they bring a stranger's text back. "Cannot
// change anything" is not the same as "cannot matter".
var readOnly = map[string]bool{
	"read_file": true,
	"glob":      true,
	"grep":      true,
	// skill reads one file from a directory found at startup, and cannot
	// reach anything read_file could not already reach.
	"skill": true,
	// ask_user changes nothing and is itself a prompt. Asking permission to
	// ask a question would be two boxes for one decision.
	"ask_user": true,
	// The plan and the budget are the agent talking about its own work. They
	// touch nothing outside the conversation.
	"todo_write":     true,
	"context_status": true,
	// Reading a background process's output reaches no further than the
	// command that produced it, which was approved when it started.
	"process_output": true,
	// Killing is the other half of starting, and the harness kills these
	// anyway when the session ends. A prompt here would ask permission to
	// tidy up — and the alternative the model reaches for without it is
	// pkill, which is far worse aimed.
	"process_kill": true,
}

// Policy decides whether a tool call proceeds, remembering what it is told.
type Policy struct {
	mu       sync.Mutex
	tools    map[string]bool // tool name -> always allowed
	commands map[string]bool // exact detail string -> always allowed

	// Prompt is called when a decision is needed. A nil Prompt denies every
	// mutating call, which is the safe default for a headless run.
	Prompt Prompt
	// YOLO skips every prompt. Set it only where the blast radius is
	// contained: a throwaway clone, a container, or the eval runner.
	YOLO bool

	// Hooks publishes the wait for a decision, for the same reason the asker
	// does: an agent sitting on an approval prompt looks exactly like one
	// thinking hard, and only one of the two will finish without you. It is
	// fired around the prompt rather than around Approve, so a call allowed
	// outright or remembered from earlier says nothing.
	Hooks *hooks.Hooks
}

// New builds a policy.
func New(prompt Prompt) *Policy {
	return &Policy{tools: map[string]bool{}, commands: map[string]bool{}, Prompt: prompt}
}

// Approve implements the agent's Approver interface. The returned error is
// text the model reads, so a denial tells it to change course rather than
// looking like a tool that broke.
func (p *Policy) Approve(ctx context.Context, tool string, input json.RawMessage) error {
	return p.ApproveAs(ctx, "", tool, input)
}

// ApproveAs is Approve on behalf of a named subagent. What is remembered is
// remembered for the whole session, not per agent: "always allow bash" is a
// statement about this working tree, and an agent id is not a trust boundary.
func (p *Policy) ApproveAs(ctx context.Context, agent, tool string, input json.RawMessage) error {
	if p == nil || p.YOLO || readOnly[tool] {
		return nil
	}
	req := Describe(tool, input)
	req.Agent = agent

	p.mu.Lock()
	allowed := p.tools[tool] || (req.Rememberable && p.commands[memoryKey(tool, req.Detail)])
	p.mu.Unlock()
	if allowed {
		return nil
	}

	if p.Prompt == nil {
		return fmt.Errorf("%s was not run: this session has no way to ask for approval, and %s can modify the working tree. Explain what you would do instead", tool, tool)
	}

	decision, err := p.decide(ctx, req)
	if err != nil {
		return err
	}
	switch decision {
	case AlwaysThisTool:
		p.mu.Lock()
		p.tools[tool] = true
		p.mu.Unlock()
		return nil
	case AlwaysThisCommand:
		if req.Rememberable {
			p.mu.Lock()
			p.commands[memoryKey(tool, req.Detail)] = true
			p.mu.Unlock()
		}
		return nil
	case Once:
		return nil
	default:
		return fmt.Errorf("the user declined to run %s. Do not retry it; either find another way or explain what you need", tool)
	}
}

// decide puts the request to the user, announcing the wait to anything
// watching from outside and announcing its end however it ends.
func (p *Policy) decide(ctx context.Context, req Request) (Decision, error) {
	p.Hooks.Fire(ctx, hooks.Payload{
		Event:   hooks.Blocked,
		Agent:   req.Agent,
		Tool:    req.Tool,
		Reason:  hooks.ReasonApproval,
		Message: blockMessage(req),
	})
	defer p.Hooks.Fire(ctx, hooks.Payload{
		Event:  hooks.Unblocked,
		Agent:  req.Agent,
		Tool:   req.Tool,
		Reason: hooks.ReasonApproval,
	})
	return p.Prompt(ctx, req)
}

// blockMessage names what is waiting to be approved. The detail is the command
// line or the path the prompt itself shows, which is what makes the difference
// between "approve something" and a decision someone can make from a phone.
func blockMessage(req Request) string {
	if req.Detail != "" {
		return req.Tool + ": " + req.Detail
	}
	return "approve " + req.Tool
}

// Remembered lists what the session has been told to always allow, for showing
// the user what they have agreed to.
func (p *Policy) Remembered() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	var out []string
	for t := range p.tools {
		out = append(out, "tool "+t)
	}
	for c := range p.commands {
		out = append(out, c)
	}
	sort.Strings(out)
	return out
}

func memoryKey(tool, detail string) string { return tool + ": " + detail }

// article picks "a" or "an" for a word, so the approval prompt does not ask
// about "a explore agent".
func article(word string) string {
	if word == "" {
		return "a"
	}
	switch word[0] {
	case 'a', 'e', 'i', 'o', 'u', 'A', 'E', 'I', 'O', 'U':
		return "an"
	}
	return "a"
}

// Describe turns a tool call into something a human can decide on in one read.
//
// It is exported because the summary is not only for a prompt box: a front end
// that shows the agent's tool calls as they happen wants the same one-line
// description, and two vocabularies for the same call would be two things to
// keep in step.
func Describe(tool string, input json.RawMessage) Request {
	var args struct {
		Command    string `json:"command"`
		Background bool   `json:"background"`
		Path       string `json:"path"`
		OldStr     string `json:"old_str"`
		NewStr     string `json:"new_str"`
		Content    string `json:"content"`
		Query      string `json:"query"`
		URL        string `json:"url"`
		Preset     string `json:"preset"`
		Task       string `json:"task"`
		Model      string `json:"model"`
		Effort     string `json:"effort"`
	}
	json.Unmarshal(input, &args)

	switch tool {
	case "bash":
		info := Analyze(args.Command)
		if args.Background {
			// Worth spelling out: agreeing to this is agreeing to something
			// that is still running after the answer comes back. Never
			// rememberable — "always allow" on a command that keeps running
			// is a standing permission to leave processes behind.
			return Request{
				Tool:         tool,
				Summary:      "start a shell command in the background, left running — " + info.Describe(),
				Detail:       args.Command,
				Rememberable: false,
			}
		}
		return Request{
			Tool:         tool,
			Summary:      "run a shell command — " + info.Describe(),
			Detail:       args.Command,
			Rememberable: !info.Opaque && info.ParseError == nil,
		}
	case "edit_file":
		return Request{
			Tool:         tool,
			Summary:      fmt.Sprintf("edit %s (%d bytes in, %d out)", args.Path, len(args.OldStr), len(args.NewStr)),
			Detail:       args.Path,
			Rememberable: true,
		}
	case "spawn_agent":
		summary := "start " + article(args.Preset) + " " + args.Preset + " agent"
		if args.Model != "" {
			summary += " on " + args.Model
		}
		if args.Effort != "" {
			summary += " at " + args.Effort + " effort"
		}
		return Request{
			Tool:    tool,
			Summary: summary,
			Detail:  args.Task,
			// The task is different every time, so remembering one exact
			// spawn would allow nothing. "Always this tool" still applies.
			Rememberable: false,
		}
	case "web_search":
		// Not rememberable, for the same reason spawn_agent's task is not: the
		// query differs every call, so remembering one exact search permits
		// nothing. "Always this tool" is the answer people give once.
		return Request{
			Tool:         tool,
			Summary:      "search the web",
			Detail:       args.Query,
			Rememberable: false,
		}
	case "web_fetch":
		return Request{
			Tool:         tool,
			Summary:      "fetch a web page",
			Detail:       args.URL,
			Rememberable: true,
		}
	case "write_file":
		return Request{
			Tool:         tool,
			Summary:      fmt.Sprintf("write %s (%d bytes, replacing whatever is there)", args.Path, len(args.Content)),
			Detail:       args.Path,
			Rememberable: true,
		}
	default:
		detail := strings.TrimSpace(string(input))
		return Request{Tool: tool, Summary: "run " + tool, Detail: detail, Rememberable: true}
	}
}