Add a codex app-server client with generated types and contract tests
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,552 @@
|
||||
package codex
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
Cwd string
|
||||
Model string
|
||||
Resume string
|
||||
GitName string
|
||||
GitEmail string
|
||||
Env []string
|
||||
Guidance string
|
||||
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 io.WriteCloser
|
||||
events chan Event
|
||||
done chan struct{}
|
||||
|
||||
writeMu sync.Mutex
|
||||
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
|
||||
|
||||
onBadLine func(line []byte, reason string)
|
||||
lastLine atomic.Int64
|
||||
}
|
||||
|
||||
func Start(ctx context.Context, 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)
|
||||
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
err = cmd.Start()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("start codex: %w", err)
|
||||
}
|
||||
|
||||
s := newSession(opts, stdin, stdout)
|
||||
s.cmd = cmd
|
||||
|
||||
err = s.open(ctx, opts)
|
||||
if err != nil {
|
||||
s.Close()
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func newSession(opts Options, stdin io.WriteCloser, stdout io.Reader) *Session {
|
||||
s := &Session{
|
||||
stdin: stdin,
|
||||
events: make(chan Event, 256),
|
||||
done: make(chan struct{}),
|
||||
transcript: opts.Transcript,
|
||||
pending: map[int64]chan rpcResult{},
|
||||
onBadLine: opts.OnBadLine,
|
||||
}
|
||||
s.lastLine.Store(time.Now().UnixNano())
|
||||
go s.readLoop(stdout)
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *Session) open(ctx context.Context, opts Options) error {
|
||||
init := InitializeResponse{}
|
||||
err := s.Call(ctx, MethodInitialize, InitializeParams{
|
||||
ClientInfo: ClientInfo{Name: "byre", Title: "byre", Version: "0.0.0"},
|
||||
}, &init)
|
||||
if err != nil {
|
||||
return fmt.Errorf("initialize: %w", err)
|
||||
}
|
||||
|
||||
s.stateMu.Lock()
|
||||
s.userAgent = init.UserAgent
|
||||
s.stateMu.Unlock()
|
||||
|
||||
err = s.notify("initialized", nil)
|
||||
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,
|
||||
}, &resp)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resume 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,
|
||||
}, &resp)
|
||||
if err != nil {
|
||||
return fmt.Errorf("start thread: %w", err)
|
||||
}
|
||||
s.setThread(resp.Thread.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
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()
|
||||
|
||||
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) Compact(ctx context.Context) error {
|
||||
return s.Call(ctx, MethodThreadCompactStart, ThreadCompactStartParams{ThreadID: s.ID()}, nil)
|
||||
}
|
||||
|
||||
func (s *Session) Call(ctx context.Context, method string, params any, result any) error {
|
||||
id := s.nextID.Add(1)
|
||||
ch := make(chan rpcResult, 1)
|
||||
|
||||
s.pendingMu.Lock()
|
||||
s.pending[id] = ch
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
defer func() {
|
||||
s.pendingMu.Lock()
|
||||
delete(s.pending, id)
|
||||
s.pendingMu.Unlock()
|
||||
}()
|
||||
|
||||
raw, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
err = s.write(message{ID: json.RawMessage(strconv.FormatInt(id, 10)), Method: method, Params: raw})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
select {
|
||||
case res := <-ch:
|
||||
if res.err != nil {
|
||||
return res.err
|
||||
}
|
||||
if result == nil {
|
||||
return nil
|
||||
}
|
||||
return json.Unmarshal(res.result, result)
|
||||
case <-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)
|
||||
}
|
||||
|
||||
func (s *Session) Close() error {
|
||||
s.writeMu.Lock()
|
||||
s.stdin.Close()
|
||||
s.writeMu.Unlock()
|
||||
|
||||
select {
|
||||
case <-s.done:
|
||||
case <-time.After(5 * time.Second):
|
||||
s.cmd.Process.Kill()
|
||||
<-s.done
|
||||
}
|
||||
|
||||
return s.cmd.Wait()
|
||||
}
|
||||
|
||||
func (s *Session) Kill() {
|
||||
s.cmd.Process.Kill()
|
||||
<-s.done
|
||||
s.cmd.Wait()
|
||||
}
|
||||
|
||||
func (s *Session) write(m message) error {
|
||||
b, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
|
||||
s.record("send", b)
|
||||
|
||||
_, err = s.stdin.Write(append(b, '\n'))
|
||||
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()}})
|
||||
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})
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
return nil, fmt.Errorf("unsupported server request %s", m.Method)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user