byre/internal/agents/codex.go

package agents

import (
	"context"
	"encoding/json"
	"log/slog"
	"os"
	"strings"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/codex"
)

var codexToolItems = map[string]bool{
	codex.ThreadItemTypeCommandExecution: true,
	codex.ThreadItemTypeFileChange:       true,
	codex.ThreadItemTypeDynamicToolCall:  true,
	codex.ThreadItemTypeMcpToolCall:      true,
	codex.ThreadItemTypeWebSearch:        true,
}

func (m *Manager) handleCodex(ctx context.Context, l *live, s *codex.Session, ev codex.Event) error {
	switch ev.Method {
	case codex.MethodItemAgentMessageDelta:
		n := codex.AgentMessageDeltaNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		m.store.Broadcast(m.event(l, agentActor, "agent.delta", map[string]any{"text": n.Delta}))
		return nil

	case codex.MethodItemStarted:
		n := codex.ItemStartedNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		if !codexToolItems[n.Item.Type] {
			return nil
		}
		name, input, err := codexToolCall(n.Item)
		if err != nil {
			return err
		}
		return m.emit(ctx, l, agentActor, "agent.tool_call", map[string]any{"id": n.Item.ID, "name": name, "input": input})

	case codex.MethodItemCompleted:
		n := codex.ItemCompletedNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		return m.codexItemCompleted(ctx, l, s, n.Item)

	case codex.MethodThreadTokenUsageUpdated:
		n := codex.ThreadTokenUsageUpdatedNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		l.contextModel = l.model
		return m.updateContext(ctx, l, int(n.TokenUsage.Last.TotalTokens), int(n.TokenUsage.ModelContextWindow))

	case codex.MethodAccountRateLimitsUpdated:
		n := codex.AccountRateLimitsUpdatedNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		m.setCodexRateLimit(n.RateLimits)
		return nil

	case codex.MethodTurnCompleted:
		n := codex.TurnCompletedNotification{}
		err := json.Unmarshal(ev.Params, &n)
		if err != nil {
			return err
		}
		return m.codexTurnCompleted(ctx, l, s, n.Turn)

	case codex.MethodError, codex.MethodWarning:
		params := map[string]any{}
		err := json.Unmarshal(ev.Params, &params)
		if err != nil {
			return err
		}
		return m.emit(ctx, l, agentActor, "agent.vendor", map[string]any{"method": ev.Method, "params": params})
	}
	return nil
}

func codexToolCall(item codex.ThreadItem) (string, any, error) {
	switch item.Type {
	case codex.ThreadItemTypeCommandExecution:
		return "Bash", map[string]any{"command": item.Command}, nil
	case codex.ThreadItemTypeFileChange:
		return "apply_patch", map[string]any{"changes": item.Changes}, nil
	case codex.ThreadItemTypeWebSearch:
		return "WebSearch", map[string]any{"query": item.Query}, nil
	case codex.ThreadItemTypeMcpToolCall:
		tool := ""
		err := json.Unmarshal(item.Tool, &tool)
		if err != nil {
			return "", nil, err
		}
		args, err := decode(item.Arguments)
		return "mcp__" + item.Server + "__" + tool, args, err
	}
	tool := ""
	err := json.Unmarshal(item.Tool, &tool)
	if err != nil {
		return "", nil, err
	}
	args, err := decode(item.Arguments)
	if item.Type == codex.ThreadItemTypeDynamicToolCall && item.Namespace != "" {
		tool = "mcp__" + item.Namespace + "__" + tool
	}
	return tool, args, err
}

func codexStatus(item codex.ThreadItem) string {
	status := ""
	json.Unmarshal(item.Status, &status)
	return status
}

func (m *Manager) codexItemCompleted(ctx context.Context, l *live, s *codex.Session, item codex.ThreadItem) error {
	switch item.Type {
	case codex.ThreadItemTypeUserMessage:
		return nil

	case codex.ThreadItemTypeAgentMessage:
		m.mu.Lock()
		l.lastText = item.Text
		m.mu.Unlock()
		return m.emit(ctx, l, agentActor, "agent.message", map[string]any{"text": item.Text})

	case codex.ThreadItemTypeReasoning:
		text := strings.TrimSpace(strings.Join(item.Summary, "\n\n"))
		if text == "" {
			return nil
		}
		return m.emit(ctx, l, agentActor, "agent.thinking", map[string]any{"text": text})
	}

	if !codexToolItems[item.Type] {
		raw, err := json.Marshal(item)
		if err != nil {
			return err
		}
		payload := map[string]any{}
		err = json.Unmarshal(raw, &payload)
		if err != nil {
			return err
		}
		return m.emit(ctx, l, agentActor, "agent.vendor", payload)
	}

	content, isError := codexToolResult(item)
	err := m.emit(ctx, l, agentActor, "agent.tool_result", map[string]any{"tool_use_id": item.ID, "is_error": isError, "content": content})
	if err != nil {
		return err
	}

	if l.reviewID != "" {
		return nil
	}
	mt := m.deliverMidTurn(ctx, l)
	if mt.text == "" {
		return nil
	}
	reply, err := s.Steer(mt.text)
	if err != nil {
		slog.Error("steer mid-turn delivery", "session", l.id, "err", err)
		m.restoreMidTurn(ctx, l, mt)
		return nil
	}
	go func() {
		ctx := context.Background()
		err := reply.Wait(ctx, nil)
		if err != nil {
			slog.Error("steer mid-turn delivery", "session", l.id, "err", err)
			m.restoreMidTurn(ctx, l, mt)
		}
	}()
	return nil
}

func codexToolResult(item codex.ThreadItem) (string, bool) {
	status := codexStatus(item)
	switch item.Type {
	case codex.ThreadItemTypeCommandExecution:
		return item.AggregatedOutput, status != "completed" || item.ExitCode != 0
	case codex.ThreadItemTypeDynamicToolCall:
		texts := []string{}
		for _, c := range item.ContentItems {
			texts = append(texts, c.Text)
		}
		return strings.Join(texts, "\n"), !item.Success
	case codex.ThreadItemTypeMcpToolCall:
		if item.Error != nil {
			return item.Error.Message, true
		}
		return string(item.Result), false
	}
	return status, status != "completed"
}

func (m *Manager) readCodexRateLimit() {
	err := os.MkdirAll(m.codexHome, 0o700)
	if err != nil {
		slog.Error("read codex rate limits", "err", err)
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	r, err := codex.ReadRateLimits(ctx, m.codexHome, []string{"CODEX_HOME=" + m.codexHome})
	if err != nil {
		slog.Error("read codex rate limits", "err", err)
		return
	}
	m.setCodexRateLimit(r)
}

func (m *Manager) setCodexRateLimit(r codex.RateLimitSnapshot) {
	m.setVendorUsage(api.Codex, codexUsage(r))
}

func codexUsage(r codex.RateLimitSnapshot) api.Usage {
	u := api.Usage{UpdatedAt: time.Now(), Windows: []api.UsageWindow{}}
	if r.RateLimitReachedType != "" {
		u.Rejected = new(string(r.RateLimitReachedType))
	}
	for _, w := range []struct {
		name string
		w    *codex.RateLimitWindow
	}{{"primary", r.Primary}, {"secondary", r.Secondary}} {
		if w.w == nil {
			continue
		}
		u.Windows = append(u.Windows, percentWindow(w.name, float64(w.w.UsedPercent)/100, w.w.ResetsAt, w.w.WindowDurationMins*60))
	}
	return u
}

func (m *Manager) codexTurnCompleted(ctx context.Context, l *live, s *codex.Session, turn codex.Turn) error {
	m.mu.Lock()
	final := l.lastText
	m.mu.Unlock()

	bounce := false
	if turn.Status == codex.TurnStatusCompleted && l.implementer() {
		var err error
		bounce, err = m.takeBounce(ctx, l, final)
		if err != nil {
			return err
		}
	}
	if bounce {
		err := m.emit(ctx, l, "byre", "user.message", map[string]any{"text": proseQuestionNudge, "mid_turn": true, "bounce": true})
		if err != nil {
			return err
		}
		m.mu.Lock()
		l.lastText = ""
		m.mu.Unlock()
		go func() {
			_, err := s.Send(context.Background(), proseQuestionNudge)
			if err != nil {
				slog.Error("send prose question nudge", "session", l.id, "err", err)
			}
		}()
		return nil
	}

	res := endResult{
		subtype: string(turn.Status),
		isError: turn.Status == codex.TurnStatusFailed,
		result:  final,
	}
	if turn.Error != nil {
		res.result = turn.Error.Message
		info := ""
		json.Unmarshal(turn.Error.CodexErrorInfo, &info)
		if info == "usageLimitExceeded" || info == "rateLimitExceeded" {
			m.mu.Lock()
			l.rateLimited = map[string]any{"status": "rejected", "rateLimitType": info}
			m.mu.Unlock()
		}
	}
	res.payload = map[string]any{
		"subtype":     res.subtype,
		"is_error":    res.isError,
		"result":      res.result,
		"vendor_uuid": turn.ID,
	}
	return m.turnEnded(ctx, l, res)
}