byre/internal/codex/session.go

package codex

import (
	"bufio"
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"g.fc.run/ai/byre/internal/mcp"
)

type Options struct {
	Cwd        string
	Model      string
	Resume     string
	ForkFrom   string
	ForkAt     string
	GitName    string
	GitEmail   string
	Env        []string
	Guidance   string
	Reminder   string
	MCPServers []*mcp.Server
	StdioMCP   []mcp.Stdio
	Transcript io.Writer
	OnBadLine  func(line []byte, reason string)
}

type Event struct {
	Method string
	Params json.RawMessage
	Read   time.Time
}

func (e Event) Decode() (any, error) {
	ctor, ok := notificationParams[e.Method]
	if !ok {
		return nil, fmt.Errorf("unknown notification %q", e.Method)
	}
	if ctor == nil {
		return nil, nil
	}
	v := ctor()
	err := json.Unmarshal(e.Params, v)
	if err != nil {
		return nil, fmt.Errorf("decode %s: %w", e.Method, err)
	}
	return v, nil
}

type Image struct {
	MediaType string
	Data      []byte
}

type RPCError struct {
	Code    int64           `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

func (e *RPCError) Error() string {
	return fmt.Sprintf("codex rpc error %d: %s", e.Code, e.Message)
}

type message struct {
	ID     json.RawMessage `json:"id,omitempty"`
	Method string          `json:"method,omitempty"`
	Params json.RawMessage `json:"params,omitempty"`
	Result json.RawMessage `json:"result,omitempty"`
	Error  *RPCError       `json:"error,omitempty"`
}

type rpcResult struct {
	result json.RawMessage
	err    error
}

type Session struct {
	cmd    *exec.Cmd
	stdin  *os.File
	events chan Event
	done   chan struct{}

	writeLock  chan struct{}
	recordMu   sync.Mutex
	transcript io.Writer

	nextID    atomic.Int64
	pendingMu sync.Mutex
	pending   map[int64]chan rpcResult

	stateMu   sync.Mutex
	threadID  string
	turnID    string
	model     string
	userAgent string

	servers  map[string]*mcp.Server
	reminder string

	onBadLine func(line []byte, reason string)
	lastLine  atomic.Int64
}

func Start(ctx context.Context, opts Options) (*Session, error) {
	s, err := spawn(opts)
	if err != nil {
		return nil, err
	}

	err = s.open(ctx, opts)
	if err != nil {
		s.Close()
		return nil, err
	}
	return s, nil
}

func ReadRateLimits(ctx context.Context, cwd string, env []string) (RateLimitSnapshot, error) {
	s, err := spawn(Options{Cwd: cwd, Env: env})
	if err != nil {
		return RateLimitSnapshot{}, err
	}
	defer s.Close()

	go func() {
		for range s.events {
		}
	}()

	err = s.initialize(ctx)
	if err != nil {
		return RateLimitSnapshot{}, err
	}

	resp := GetAccountRateLimitsResponse{}
	err = s.Call(ctx, MethodAccountRateLimitsRead, GetAccountRateLimitsParams{ExcludeResetCreditDetails: true}, &resp)
	if err != nil {
		return RateLimitSnapshot{}, fmt.Errorf("read rate limits: %w", err)
	}
	return resp.RateLimits, nil
}

func spawn(opts Options) (*Session, error) {
	cmd := exec.Command("codex", "app-server", "-c", "check_for_update_on_startup=false")
	cmd.Dir = opts.Cwd
	cmd.Stderr = os.Stderr
	cmd.Env = buildEnv(opts)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	stdinR, stdin, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	cmd.Stdin = stdinR

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		stdinR.Close()
		stdin.Close()
		return nil, err
	}

	err = cmd.Start()
	stdinR.Close()
	if err != nil {
		stdin.Close()
		return nil, fmt.Errorf("start codex: %w", err)
	}

	s := newSession(opts, stdin, stdout)
	s.cmd = cmd
	return s, nil
}

func newSession(opts Options, stdin *os.File, stdout io.Reader) *Session {
	s := &Session{
		stdin:      stdin,
		events:     make(chan Event, 256),
		writeLock:  make(chan struct{}, 1),
		done:       make(chan struct{}),
		transcript: opts.Transcript,
		pending:    map[int64]chan rpcResult{},
		servers:    map[string]*mcp.Server{},
		reminder:   opts.Reminder,
		onBadLine:  opts.OnBadLine,
	}
	for _, srv := range opts.MCPServers {
		s.servers[srv.Name] = srv
	}
	s.lastLine.Store(time.Now().UnixNano())
	go s.readLoop(stdout)
	return s
}

func (s *Session) open(ctx context.Context, opts Options) error {
	err := s.initialize(ctx)
	if err != nil {
		return err
	}

	if opts.Resume != "" {
		resp := ThreadResumeResponse{}
		err = s.Call(ctx, MethodThreadResume, ThreadResumeParams{
			ThreadID:              opts.Resume,
			Cwd:                   opts.Cwd,
			Model:                 opts.Model,
			ApprovalPolicy:        json.RawMessage(`"never"`),
			Sandbox:               SandboxModeDangerFullAccess,
			DeveloperInstructions: opts.Guidance,
			Config:                stdioConfig(opts.StdioMCP),
		}, &resp)
		if err != nil {
			return fmt.Errorf("resume thread: %w", err)
		}
		s.setThread(resp.Thread.ID)
		return nil
	}

	if opts.ForkFrom != "" {
		resp := ThreadForkResponse{}
		err = s.Call(ctx, MethodThreadFork, ThreadForkParams{
			ThreadID:              opts.ForkFrom,
			LastTurnID:            opts.ForkAt,
			Cwd:                   opts.Cwd,
			Model:                 opts.Model,
			ApprovalPolicy:        json.RawMessage(`"never"`),
			Sandbox:               SandboxModeDangerFullAccess,
			DeveloperInstructions: opts.Guidance,
			Config:                stdioConfig(opts.StdioMCP),
		}, &resp)
		if err != nil {
			return fmt.Errorf("fork thread: %w", err)
		}
		s.setThread(resp.Thread.ID)
		return nil
	}

	resp := ThreadStartResponse{}
	err = s.Call(ctx, MethodThreadStart, ThreadStartParams{
		Cwd:                   opts.Cwd,
		Model:                 opts.Model,
		ApprovalPolicy:        json.RawMessage(`"never"`),
		Sandbox:               SandboxModeDangerFullAccess,
		DeveloperInstructions: opts.Guidance,
		DynamicTools:          dynamicTools(opts.MCPServers),
		Config:                stdioConfig(opts.StdioMCP),
	}, &resp)
	if err != nil {
		return fmt.Errorf("start thread: %w", err)
	}
	s.setThread(resp.Thread.ID)
	return nil
}

func (s *Session) initialize(ctx context.Context) error {
	init := InitializeResponse{}
	err := s.Call(ctx, MethodInitialize, InitializeParams{
		ClientInfo:   ClientInfo{Name: "byre", Title: "byre", Version: "0.0.0"},
		Capabilities: &InitializeCapabilities{ExperimentalApi: true},
	}, &init)
	if err != nil {
		return fmt.Errorf("initialize: %w", err)
	}

	s.stateMu.Lock()
	s.userAgent = init.UserAgent
	s.stateMu.Unlock()

	return s.notify("initialized", nil)
}

func stdioConfig(servers []mcp.Stdio) map[string]any {
	if len(servers) == 0 {
		return nil
	}
	cfg := map[string]any{}
	for _, srv := range servers {
		cfg[srv.Name] = map[string]any{"command": srv.Command, "args": srv.Args}
	}
	return map[string]any{"mcp_servers": cfg}
}

func dynamicTools(servers []*mcp.Server) []DynamicToolSpec {
	specs := []DynamicToolSpec{}
	for _, srv := range servers {
		tools := []DynamicToolNamespaceTool{}
		for _, t := range srv.Tools {
			tools = append(tools, DynamicToolNamespaceTool{
				Type:        DynamicToolSpecTypeFunction,
				Name:        t.Name,
				Description: t.Description,
				InputSchema: t.InputSchema,
			})
		}
		specs = append(specs, DynamicToolSpec{
			Type:        DynamicToolSpecTypeNamespace,
			Name:        srv.Name,
			Description: srv.Name + "'s tools",
			Tools:       tools,
		})
	}
	return specs
}

func (s *Session) setThread(id string) {
	s.stateMu.Lock()
	s.threadID = id
	s.stateMu.Unlock()
}

func buildEnv(opts Options) []string {
	env := append(os.Environ(),
		"GIT_AUTHOR_NAME="+opts.GitName,
		"GIT_AUTHOR_EMAIL="+opts.GitEmail,
		"GIT_COMMITTER_NAME="+opts.GitName,
		"GIT_COMMITTER_EMAIL="+opts.GitEmail,
	)
	return append(env, opts.Env...)
}

func (s *Session) Events() <-chan Event {
	return s.events
}

func (s *Session) Done() <-chan struct{} {
	return s.done
}

func (s *Session) Pid() int {
	return s.cmd.Process.Pid
}

func (s *Session) LastLine() time.Time {
	return time.Unix(0, s.lastLine.Load())
}

func (s *Session) Exited() bool {
	select {
	case <-s.done:
		return true
	default:
		return false
	}
}

func (s *Session) ID() string {
	s.stateMu.Lock()
	defer s.stateMu.Unlock()
	return s.threadID
}

func (s *Session) TurnID() string {
	s.stateMu.Lock()
	defer s.stateMu.Unlock()
	return s.turnID
}

func (s *Session) UserAgent() string {
	s.stateMu.Lock()
	defer s.stateMu.Unlock()
	return s.userAgent
}

func (s *Session) SetModel(model string) {
	s.stateMu.Lock()
	s.model = model
	s.stateMu.Unlock()
}

func (s *Session) Send(ctx context.Context, text string, images ...Image) (Turn, error) {
	input := []UserInput{}
	for _, img := range images {
		input = append(input, UserInput{
			Type: "image",
			URL:  "data:" + img.MediaType + ";base64," + base64.StdEncoding.EncodeToString(img.Data),
		})
	}
	if text != "" {
		input = append(input, UserInput{Type: "text", Text: text})
	}

	s.stateMu.Lock()
	params := TurnStartParams{ThreadID: s.threadID, Input: input, Model: s.model}
	s.model = ""
	s.stateMu.Unlock()
	if s.reminder != "" {
		params.AdditionalContext = map[string]AdditionalContextEntry{
			"byre": {Kind: AdditionalContextKindApplication, Value: s.reminder},
		}
	}

	resp := TurnStartResponse{}
	err := s.Call(ctx, MethodTurnStart, params, &resp)
	if err != nil {
		return Turn{}, err
	}

	s.stateMu.Lock()
	if s.turnID == "" {
		s.turnID = resp.Turn.ID
	}
	s.stateMu.Unlock()
	return resp.Turn, nil
}

func (s *Session) Interrupt(ctx context.Context) error {
	s.stateMu.Lock()
	params := TurnInterruptParams{ThreadID: s.threadID, TurnID: s.turnID}
	s.stateMu.Unlock()

	if params.TurnID == "" {
		return errors.New("no turn in progress")
	}
	return s.Call(ctx, MethodTurnInterrupt, params, nil)
}

func (s *Session) Steer(text string) (*Reply, error) {
	s.stateMu.Lock()
	params := TurnSteerParams{ThreadID: s.threadID, ExpectedTurnID: s.turnID, Input: []UserInput{{Type: UserInputTypeText, Text: text}}}
	s.stateMu.Unlock()

	if params.ExpectedTurnID == "" {
		return nil, errors.New("no turn in progress")
	}
	return s.request(context.Background(), MethodTurnSteer, params)
}

func (s *Session) Compact(ctx context.Context) error {
	return s.Call(ctx, MethodThreadCompactStart, ThreadCompactStartParams{ThreadID: s.ID()}, nil)
}

const rpcTimeout = time.Minute

type Reply struct {
	s        *Session
	id       int64
	ch       chan rpcResult
	deadline time.Time
}

func (s *Session) Call(ctx context.Context, method string, params any, result any) error {
	r, err := s.request(ctx, method, params)
	if err != nil {
		return err
	}
	return r.Wait(ctx, result)
}

func rpcDeadline(ctx context.Context) time.Time {
	deadline := time.Now().Add(rpcTimeout)
	d, ok := ctx.Deadline()
	if ok && d.Before(deadline) {
		return d
	}
	return deadline
}

func (s *Session) request(ctx context.Context, method string, params any) (*Reply, error) {
	raw, err := json.Marshal(params)
	if err != nil {
		return nil, err
	}

	r := &Reply{s: s, id: s.nextID.Add(1), ch: make(chan rpcResult, 1), deadline: rpcDeadline(ctx)}

	s.pendingMu.Lock()
	s.pending[r.id] = r.ch
	s.pendingMu.Unlock()

	err = s.write(message{ID: json.RawMessage(strconv.FormatInt(r.id, 10)), Method: method, Params: raw}, r.deadline)
	if err != nil {
		r.forget()
		return nil, err
	}
	return r, nil
}

func (r *Reply) forget() {
	r.s.pendingMu.Lock()
	delete(r.s.pending, r.id)
	r.s.pendingMu.Unlock()
}

func (r *Reply) Wait(ctx context.Context, result any) error {
	defer r.forget()

	ctx, cancel := context.WithDeadline(ctx, r.deadline)
	defer cancel()

	select {
	case res := <-r.ch:
		if res.err != nil {
			return res.err
		}
		if result == nil {
			return nil
		}
		return json.Unmarshal(res.result, result)
	case <-r.s.done:
		return errors.New("codex exited")
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *Session) notify(method string, params any) error {
	m := message{Method: method}
	if params != nil {
		raw, err := json.Marshal(params)
		if err != nil {
			return err
		}
		m.Params = raw
	}
	return s.write(m, time.Now().Add(rpcTimeout))
}

func (s *Session) Close() error {
	s.writeLock <- struct{}{}
	s.stdin.Close()
	<-s.writeLock

	select {
	case <-s.done:
	case <-time.After(5 * time.Second):
		syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
		<-s.done
	}

	return s.cmd.Wait()
}

func (s *Session) Kill() {
	syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
	<-s.done
	s.cmd.Wait()
}

func (s *Session) write(m message, deadline time.Time) error {
	b, err := json.Marshal(m)
	if err != nil {
		return err
	}

	wait := time.NewTimer(time.Until(deadline))
	defer wait.Stop()
	select {
	case s.writeLock <- struct{}{}:
	case <-wait.C:
		return os.ErrDeadlineExceeded
	}
	defer func() { <-s.writeLock }()

	if !time.Now().Before(deadline) {
		return os.ErrDeadlineExceeded
	}

	s.record("send", b)

	err = s.stdin.SetWriteDeadline(deadline)
	if err != nil {
		return err
	}
	n, err := s.stdin.Write(append(b, '\n'))
	if n > 0 && errors.Is(err, os.ErrDeadlineExceeded) {
		slog.Error("codex stopped reading its stdin mid-line; killing it", "pid", s.cmd.Process.Pid, "written", n, "len", len(b)+1)
		syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
	}
	return err
}

var emailField = regexp.MustCompile(`"email":"[^"]*"`)

func (s *Session) record(dir string, line []byte) {
	if s.transcript == nil {
		return
	}

	line = emailField.ReplaceAll(line, []byte(`"email":"[redacted]"`))

	b, err := json.Marshal(map[string]any{"dir": dir, "msg": json.RawMessage(line)})
	if err != nil {
		slog.Error("record transcript", "err", err)
		return
	}

	s.recordMu.Lock()
	defer s.recordMu.Unlock()
	s.transcript.Write(append(b, '\n'))
}

func (s *Session) readLoop(stdout io.Reader) {
	defer close(s.done)
	defer close(s.events)

	r := bufio.NewReader(stdout)
	for {
		line, err := r.ReadBytes('\n')
		line = bytes.TrimRight(line, "\r\n")
		if len(line) > 0 {
			s.lastLine.Store(time.Now().UnixNano())
		}
		if len(line) > 0 && line[0] == '{' {
			s.handleLine(line)
		} else if len(line) > 0 {
			s.badLine(line, "not a JSON object")
		}
		if err != nil {
			if err != io.EOF {
				slog.Error("read codex stdout", "err", err)
			}
			return
		}
	}
}

func (s *Session) handleLine(line []byte) {
	read := time.Now()
	m := message{}
	err := json.Unmarshal(line, &m)
	if err != nil {
		s.badLine(line, err.Error())
		return
	}

	s.record("recv", line)

	if m.Method == "" {
		s.handleResponse(line, m)
		return
	}

	if len(m.ID) > 0 {
		go s.handleServerRequest(m)
		return
	}

	if _, ok := notificationParams[m.Method]; !ok {
		s.badLine(line, fmt.Sprintf("unrecognised method %q", m.Method))
	}

	s.trackTurn(m)
	s.events <- Event{Method: m.Method, Params: append(json.RawMessage{}, m.Params...), Read: read}
}

func (s *Session) trackTurn(m message) {
	switch m.Method {
	case MethodTurnStarted:
		n := TurnStartedNotification{}
		if json.Unmarshal(m.Params, &n) == nil {
			s.stateMu.Lock()
			s.turnID = n.Turn.ID
			s.stateMu.Unlock()
		}
	case MethodTurnCompleted:
		n := TurnCompletedNotification{}
		if json.Unmarshal(m.Params, &n) == nil {
			s.stateMu.Lock()
			if s.turnID == n.Turn.ID {
				s.turnID = ""
			}
			s.stateMu.Unlock()
		}
	}
}

func (s *Session) handleResponse(line []byte, m message) {
	id, err := strconv.ParseInt(string(m.ID), 10, 64)
	if err != nil {
		s.badLine(line, "response with a non-numeric id")
		return
	}

	s.pendingMu.Lock()
	ch := s.pending[id]
	s.pendingMu.Unlock()

	if ch == nil {
		slog.Warn("codex response for unknown request", "id", id)
		return
	}

	if m.Error != nil {
		ch <- rpcResult{err: m.Error}
		return
	}
	ch <- rpcResult{result: m.Result}
}

func (s *Session) handleServerRequest(m message) {
	result, err := s.answerServerRequest(m)
	if err != nil {
		slog.Error("codex server request failed", "method", m.Method, "err", err)
		s.write(message{ID: m.ID, Error: &RPCError{Code: -32601, Message: err.Error()}}, time.Now().Add(rpcTimeout))
		return
	}

	raw, err := json.Marshal(result)
	if err != nil {
		slog.Error("encode codex server response", "method", m.Method, "err", err)
		return
	}
	s.write(message{ID: m.ID, Result: raw}, time.Now().Add(rpcTimeout))
}

func (s *Session) answerServerRequest(m message) (any, error) {
	switch m.Method {
	case MethodItemCommandExecutionRequestApproval:
		return CommandExecutionRequestApprovalResponse{Decision: json.RawMessage(`"accept"`)}, nil
	case MethodItemFileChangeRequestApproval:
		return FileChangeRequestApprovalResponse{Decision: "accept"}, nil
	case MethodItemToolCall:
		return s.callTool(m.Params)
	}
	return nil, fmt.Errorf("unsupported server request %s", m.Method)
}

func (s *Session) callTool(raw json.RawMessage) (any, error) {
	p := DynamicToolCallParams{}
	err := json.Unmarshal(raw, &p)
	if err != nil {
		return nil, err
	}

	srv := s.servers[p.Namespace]
	if srv == nil {
		return nil, fmt.Errorf("unknown tool namespace %q", p.Namespace)
	}
	t, ok := srv.Tool(p.Tool)
	if !ok {
		return nil, fmt.Errorf("unknown tool %s in %s", p.Tool, p.Namespace)
	}

	res := t.Call(context.Background(), p.Arguments)
	items := []DynamicToolCallOutputContentItem{}
	for _, c := range res.Content {
		if c.Type == "image" {
			items = append(items, DynamicToolCallOutputContentItem{Type: DynamicToolCallOutputContentItemTypeInputImage, ImageURL: "data:" + c.MimeType + ";base64," + c.Data})
			continue
		}
		items = append(items, DynamicToolCallOutputContentItem{Type: DynamicToolCallOutputContentItemTypeInputText, Text: c.Text})
	}
	return DynamicToolCallResponse{Success: !res.IsError, ContentItems: items}, nil
}

func (s *Session) badLine(line []byte, reason string) {
	slog.Warn("bad codex line", "reason", reason, "line", string(line))
	if s.onBadLine != nil {
		s.onBadLine(line, reason)
	}
}