Vectrum
Sign In
157 lines · 4.7 KB · Go
package serve

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"sync"

	"codey/internal/acp"
	"codey/internal/agent"
	"codey/internal/approval"
)

// session is one conversation. A client may run several at once against the
// same process — T3 Code opens one per thread — so everything a turn touches
// lives here rather than on the server.
type session struct {
	id   string
	root string

	// state is the conversation, and it outlives any one model choice:
	// selecting another model rebuilds the agent around this same state, so
	// the history survives the switch.
	state  *agent.State
	policy *approval.Policy

	// mu serializes turns. A second prompt while one is running waits for it,
	// which is the same thing the TUI does and the only thing the agent's
	// history can survive.
	mu    sync.Mutex
	agent *agent.Agent
	model string
	emit  *emitter

	// cmu guards the cancellation of the turn in flight. It is separate from
	// mu because a cancel arrives while mu is held by the turn it cancels.
	cmu       sync.Mutex
	cancel    context.CancelFunc
	cancelled bool
}

// newSessionID names a session. It is ours to choose and the client only
// echoes it back, so the only requirements are that it is unique within the
// process and recognisable in a log.
func newSessionID() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		// crypto/rand does not fail in practice, and a session id is not a
		// secret: a readable fallback beats failing to open a conversation.
		return "codey-fallback"
	}
	return "codey-" + hex.EncodeToString(b[:])
}

// begin makes a turn cancellable and clears any cancellation left over from
// the last one. A stop button pressed after a turn ended must not cancel the
// next one.
func (s *session) begin(parent context.Context) context.Context {
	ctx, cancel := context.WithCancel(parent)
	s.cmu.Lock()
	s.cancel, s.cancelled = cancel, false
	s.cmu.Unlock()
	return ctx
}

func (s *session) end() {
	s.cmu.Lock()
	cancel := s.cancel
	s.cancel = nil
	s.cmu.Unlock()
	if cancel != nil {
		cancel()
	}
}

// stop cancels the turn in flight, if there is one. It reports nothing: the
// acknowledgement the client gets is the prompt coming back cancelled.
func (s *session) stop() {
	s.cmu.Lock()
	cancel := s.cancel
	s.cancelled = true
	s.cmu.Unlock()
	if cancel != nil {
		cancel()
	}
}

func (s *session) wasCancelled() bool {
	s.cmu.Lock()
	defer s.cmu.Unlock()
	return s.cancelled
}

// promptFor builds the approval prompt for this session: a permission request
// over the wire, and whatever the client answers.
//
// The tool call id is the one the emitter just announced. The agent emits
// tool_start before it asks for approval, and dispatch runs one call at a
// time, so "the call currently in flight" is exactly the one being decided —
// and attaching the prompt to that row is what lets a client show the question
// beside the command it is about.
func (s *session) promptFor(conn *acp.Conn, mode PermissionMode) approval.Prompt {
	return func(ctx context.Context, req approval.Request) (approval.Decision, error) {
		if mode.allows(req) {
			return approval.Once, nil
		}

		params := acp.RequestPermissionParams{
			SessionID: s.id,
			ToolCall: acp.PermissionToolCall{
				ToolCallID: s.emit.currentCall(),
				Title:      title(req),
				Kind:       toolKind(req.Tool),
				Status:     acp.StatusPending,
			},
			Options: options(req),
		}

		var result acp.RequestPermissionResult
		if err := conn.Request(ctx, acp.MethodRequestPermission, params, &result); err != nil {
			return approval.Deny, err
		}

		if result.Outcome.Outcome != acp.OutcomeSelected {
			return approval.Deny, context.Canceled
		}
		switch result.Outcome.OptionID {
		case acp.OptAllowAlways:
			// "Always" means the tool from here on when the call itself cannot
			// be pinned down — an opaque shell command is not a thing to
			// remember by name.
			if req.Rememberable {
				return approval.AlwaysThisCommand, nil
			}
			return approval.AlwaysThisTool, nil
		case acp.OptAllowOnce:
			return approval.Once, nil
		default:
			return approval.Deny, nil
		}
	}
}

// options are the three answers offered. The ids are fixed literals rather
// than anything derived: a client that replies with an id it chose itself —
// which is what T3 Code's Cursor adapter does — must find one of these.
func options(req approval.Request) []acp.PermissionOption {
	always := "Always allow this command"
	if !req.Rememberable {
		always = "Always allow " + req.Tool
	}
	return []acp.PermissionOption{
		{OptionID: acp.OptAllowOnce, Name: "Allow once", Kind: acp.KindAllowOnce},
		{OptionID: acp.OptAllowAlways, Name: always, Kind: acp.KindAllowAlways},
		{OptionID: acp.OptRejectOnce, Name: "Reject", Kind: acp.KindRejectOnce},
	}
}