chenhg5-cc-connect/agent/codex/session.go

package codex

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"

	"github.com/chenhg5/cc-connect/core"
)

// codexSession manages a multi-turn Codex conversation.
// First Send() uses `codex exec`, subsequent ones use `codex exec resume <threadID>`.
type codexSession struct {
	workDir       string
	model         string
	effort        string
	mode          string
	baseURL       string // provider base URL; passed as -c openai_base_url=<url>
	modelProvider string // Codex model_provider name; passed as -c model_provider=<name>
	cliBin        string   // CLI binary, default "codex"
	cliExtraArgs  []string // extra args from cli_path, prepended before exec args
	extraEnv      []string
	events        chan core.Event
	threadID  atomic.Value // stores string — Codex thread_id
	ctx       context.Context
	cancel    context.CancelFunc
	wg        sync.WaitGroup
	alive     atomic.Bool
	closeOnce sync.Once
	cmdMu     sync.Mutex
	cmds      map[*exec.Cmd]struct{}

	pendingMsgs []string // buffered agent_message texts awaiting classification

	runtimeCfgMu       sync.Mutex
	runtimeCfgModel    string
	runtimeCfgEffort   string
	runtimeCfgFetched  time.Time
	runtimeCfgFetchErr error

	contextMu    sync.RWMutex
	contextUsage *core.ContextUsage
	sessionFile  string
}

var codexSessionCloseTimeout = 8 * time.Second
var codexSessionForceKillWait = 2 * time.Second
var codexRuntimeConfigCacheTTL = 5 * time.Second
var codexRuntimeConfigTimeout = 1500 * time.Millisecond
var codexContextUsageRetryDelay = 50 * time.Millisecond
var codexContextUsageRetryCount = 4

func newCodexSession(ctx context.Context, cliBin string, cliExtraArgs []string, workDir, model, effort, mode, resumeID, baseURL string, extraEnv []string, modelProvider string) (*codexSession, error) {
	sessionCtx, cancel := context.WithCancel(ctx)

	cs := &codexSession{
		workDir:       workDir,
		model:         model,
		effort:        effort,
		mode:          mode,
		baseURL:       baseURL,
		modelProvider: modelProvider,
		cliBin:        cliBin,
		cliExtraArgs:  cliExtraArgs,
		extraEnv:      extraEnv,
		events:        make(chan core.Event, 64),
		ctx:           sessionCtx,
		cancel:        cancel,
		cmds:          make(map[*exec.Cmd]struct{}),
	}
	cs.alive.Store(true)

	if resumeID != "" && resumeID != core.ContinueSession {
		cs.threadID.Store(resumeID)
	}

	return cs, nil
}

// Send launches a codex subprocess.
// If a threadID exists (from a prior turn or resume), uses `codex exec resume <id> <prompt>`.
// Otherwise uses `codex exec <prompt>` to start a new conversation.
func (cs *codexSession) Send(prompt string, images []core.ImageAttachment, files []core.FileAttachment) error {
	if len(files) > 0 {
		filePaths := core.SaveFilesToDisk(cs.workDir, files)
		prompt = core.AppendFileRefs(prompt, filePaths)
	}
	if !cs.alive.Load() {
		return fmt.Errorf("session is closed")
	}

	prompt, imagePaths, err := cs.stageImages(prompt, images)
	if err != nil {
		return err
	}

	isResume := cs.CurrentSessionID() != ""
	args := cs.buildExecArgs(prompt, imagePaths)
	if len(cs.cliExtraArgs) > 0 {
		args = append(append([]string{}, cs.cliExtraArgs...), args...)
	}

	bin := cs.cliBin
	if bin == "" {
		bin = "codex"
	}

	slog.Debug("codexSession: launching", "resume", isResume, "args", core.RedactArgs(args))

	cmd := exec.CommandContext(cs.ctx, bin, args...)
	cmd.Dir = cs.workDir
	prepareCmdForKill(cmd)
	if len(cs.extraEnv) > 0 {
		cmd.Env = core.MergeEnv(os.Environ(), cs.extraEnv)
	}
	cmd.Stdin = strings.NewReader(prompt)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("codexSession: stdout pipe: %w", err)
	}

	var stderrBuf bytes.Buffer
	cmd.Stderr = &stderrBuf

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("codexSession: start: %w", err)
	}
	cs.addCmd(cmd)

	cs.wg.Add(1)
	go cs.readLoop(cmd, stdout, &stderrBuf)

	return nil
}

func (cs *codexSession) stageImages(prompt string, images []core.ImageAttachment) (string, []string, error) {
	if len(images) == 0 {
		return prompt, nil, nil
	}

	imgDir := filepath.Join(cs.workDir, ".cc-connect", "images")
	if err := os.MkdirAll(imgDir, 0o755); err != nil {
		return "", nil, fmt.Errorf("codexSession: create image dir: %w", err)
	}

	imagePaths := make([]string, 0, len(images))
	for i, img := range images {
		ext := codexImageExt(img.MimeType)
		fname := fmt.Sprintf("img_%d_%d%s", time.Now().UnixMilli(), i, ext)
		fpath := filepath.Join(imgDir, fname)
		if err := os.WriteFile(fpath, img.Data, 0o644); err != nil {
			return "", nil, fmt.Errorf("codexSession: save image: %w", err)
		}
		imagePaths = append(imagePaths, fpath)
	}

	if strings.TrimSpace(prompt) == "" {
		prompt = "Please analyze the attached image(s)."
	}

	return prompt, imagePaths, nil
}

func (cs *codexSession) buildExecArgs(prompt string, imagePaths []string) []string {
	tid := cs.CurrentSessionID()
	isResume := tid != ""

	var args []string
	if isResume {
		// For resume: codex exec resume ... <thread_id> [--image ...] --json --cd <dir> <prompt>
		// The codex CLI requires --json after the thread_id positional argument.
		args = []string{"exec", "resume", "--skip-git-repo-check"}
	} else {
		args = []string{"exec", "--skip-git-repo-check"}
	}

	switch cs.mode {
	case "auto-edit", "full-auto":
		args = append(args, "--full-auto")
	case "yolo":
		args = append(args, "--dangerously-bypass-approvals-and-sandbox")
	}

	if cs.model != "" {
		args = append(args, "--model", cs.model)
	}
	if cs.modelProvider != "" {
		args = append(args, "-c", fmt.Sprintf("model_provider=%q", cs.modelProvider))
	}
	if cs.baseURL != "" {
		args = append(args, "-c", fmt.Sprintf("openai_base_url=%q", cs.baseURL))
	}
	if cs.effort != "" {
		args = append(args, "-c", fmt.Sprintf("model_reasoning_effort=%q", cs.effort))
	}

	if isResume {
		args = append(args, tid)
		for _, imagePath := range imagePaths {
			args = append(args, "--image", imagePath)
		}
		// codex exec resume does not support --cd; cmd.Dir handles cwd instead.
		// Use stdin ("-") so multiline prompts are preserved reliably on Windows.
		args = append(args, "--json", "-")
	} else {
		for _, imagePath := range imagePaths {
			args = append(args, "--image", imagePath)
		}
		args = append(args, "--json", "--cd", cs.workDir, "-")
	}
	return args
}

func codexImageExt(mime string) string {
	switch mime {
	case "image/jpeg":
		return ".jpg"
	case "image/gif":
		return ".gif"
	case "image/webp":
		return ".webp"
	default:
		return ".png"
	}
}

func (cs *codexSession) readLoop(cmd *exec.Cmd, stdout io.ReadCloser, stderrBuf *bytes.Buffer) {
	defer cs.wg.Done()
	defer func() {
		defer cs.removeCmd(cmd)
		if err := cmd.Wait(); err != nil {
			stderrMsg := strings.TrimSpace(stderrBuf.String())
			if stderrMsg != "" {
				slog.Error("codexSession: process failed", "error", err, "stderr", stderrMsg)
				evt := core.Event{Type: core.EventError, Error: fmt.Errorf("%s", stderrMsg)}
				select {
				case cs.events <- evt:
				case <-cs.ctx.Done():
					return
				}
			}
		}
		if tid := cs.CurrentSessionID(); tid != "" {
			patchSessionSource(tid, getenvFromList(cs.extraEnv, "CODEX_HOME"))
		}
	}()

	if err := readJSONLines(stdout, func(line []byte) error {
		lineText := string(line)
		if lineText == "" {
			return nil
		}

		slog.Debug("codexSession: raw", "line", truncate(lineText, 500))

		var raw map[string]any
		if err := json.Unmarshal(line, &raw); err != nil {
			slog.Debug("codexSession: non-JSON line", "line", lineText)
			return nil
		}

		cs.handleEvent(raw)
		return nil
	}); err != nil {
		slog.Error("codexSession: read stdout error", "error", err)
		evt := core.Event{Type: core.EventError, Error: fmt.Errorf("read stdout: %w", err)}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}
	}
}

func readJSONLines(r io.Reader, handle func([]byte) error) error {
	reader := bufio.NewReader(r)

	for {
		line, err := reader.ReadBytes('\n')
		if errors.Is(err, io.EOF) && len(line) == 0 {
			return nil
		}
		if err != nil && !errors.Is(err, io.EOF) {
			return err
		}

		line = bytes.TrimRight(line, "\r\n")
		if len(line) > 0 {
			if err := handle(line); err != nil {
				return err
			}
		}

		if errors.Is(err, io.EOF) {
			return nil
		}
	}
}

func (cs *codexSession) handleEvent(raw map[string]any) {
	eventType, _ := raw["type"].(string)

	switch eventType {
	case "thread.started":
		if tid, ok := raw["thread_id"].(string); ok {
			cs.threadID.Store(tid)
			cs.contextMu.Lock()
			cs.sessionFile = ""
			cs.contextUsage = nil
			cs.contextMu.Unlock()
			slog.Debug("codexSession: thread started", "thread_id", tid)
		}

	case "turn.started":
		cs.pendingMsgs = cs.pendingMsgs[:0]
		cs.contextMu.Lock()
		cs.contextUsage = nil
		cs.contextMu.Unlock()
		slog.Debug("codexSession: turn started")

	case "item.started":
		cs.handleItemStarted(raw)

	case "item.completed":
		cs.handleItemCompleted(raw)

	case "turn.completed":
		cs.refreshContextUsageFromRollout()
		cs.flushPendingAsText()
		evt := core.Event{Type: core.EventResult, SessionID: cs.CurrentSessionID(), Done: true}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}

	case "turn.failed":
		errMsg := ""
		if errObj, ok := raw["error"].(map[string]any); ok {
			errMsg, _ = errObj["message"].(string)
		}
		if errMsg == "" {
			errMsg = "turn failed (no details)"
		}
		slog.Warn("codexSession: turn failed", "error", errMsg)
		evt := core.Event{Type: core.EventError, Error: fmt.Errorf("%s", errMsg)}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}

	case "error":
		msg, _ := raw["message"].(string)
		if strings.Contains(msg, "Reconnecting") || strings.Contains(msg, "Falling back") {
			slog.Debug("codexSession: transient error", "message", msg)
		} else {
			slog.Warn("codexSession: error event", "message", msg)
		}

	default:
		slog.Debug("codexSession: unhandled event type", "type", eventType)
	}
}

// flushPendingAsThinking emits all buffered agent_messages as EventThinking.
func (cs *codexSession) flushPendingAsThinking() {
	if cs.ctx.Err() != nil {
		return
	}
	for _, text := range cs.pendingMsgs {
		if cs.ctx.Err() != nil {
			return
		}
		evt := core.Event{Type: core.EventThinking, Content: text}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}
	}
	cs.pendingMsgs = cs.pendingMsgs[:0]
}

// flushPendingAsText emits all buffered agent_messages as EventText (final response).
func (cs *codexSession) flushPendingAsText() {
	if cs.ctx.Err() != nil {
		return
	}
	for _, text := range cs.pendingMsgs {
		if cs.ctx.Err() != nil {
			return
		}
		evt := core.Event{Type: core.EventText, Content: text}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}
	}
	cs.pendingMsgs = cs.pendingMsgs[:0]
}

var codexToolNames = map[string]string{
	"web_search":       "WebSearch",
	"file_search":      "FileSearch",
	"code_interpreter": "CodeInterpreter",
	"computer_use":     "ComputerUse",
	"mcp_tool":         "MCP",
}

func (cs *codexSession) handleItemStarted(raw map[string]any) {
	item, ok := raw["item"].(map[string]any)
	if !ok {
		slog.Debug("codexSession: item.started missing item field")
		return
	}
	itemType, _ := item["type"].(string)
	slog.Debug("codexSession: item.started", "item_type", itemType)

	if itemType == "agent_message" || itemType == "message" || itemType == "reasoning" {
		return
	}

	// Any non-message item is a tool use; flush pending messages as thinking first.
	cs.flushPendingAsThinking()

	switch itemType {
	case "command_execution":
		command, _ := item["command"].(string)
		evt := core.Event{Type: core.EventToolUse, ToolName: "Bash", ToolInput: command}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}
	case "function_call":
		name, _ := item["name"].(string)
		args, _ := item["arguments"].(string)
		evt := core.Event{Type: core.EventToolUse, ToolName: name, ToolInput: args}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}
	}
	// Other tool types (web_search etc.) have empty fields at start;
	// their EventToolUse is emitted from handleItemCompleted instead.
}

func (cs *codexSession) handleItemCompleted(raw map[string]any) {
	item, ok := raw["item"].(map[string]any)
	if !ok {
		slog.Debug("codexSession: item.completed missing item field")
		return
	}
	itemType, _ := item["type"].(string)
	slog.Debug("codexSession: item.completed", "item_type", itemType)

	switch itemType {
	case "reasoning":
		text := extractItemText(item, "summary", "summary_text")
		if text != "" {
			evt := core.Event{Type: core.EventThinking, Content: text}
			select {
			case cs.events <- evt:
			case <-cs.ctx.Done():
				return
			}
		}

	case "agent_message", "message":
		text := extractItemText(item, "content", "output_text")
		if text != "" {
			cs.pendingMsgs = append(cs.pendingMsgs, text)
		}

	case "command_execution":
		command, _ := item["command"].(string)
		status, _ := item["status"].(string)
		output, _ := item["aggregated_output"].(string)
		exitCode, _ := item["exit_code"].(float64)
		code := int(exitCode)
		success := codexToolSuccess(status, &code)

		slog.Debug("codexSession: command completed",
			"command", truncate(command, 100),
			"status", status,
			"exit_code", code,
			"output_len", len(output),
		)
		evt := core.Event{
			Type:         core.EventToolResult,
			ToolName:     "Bash",
			ToolResult:   truncate(strings.TrimSpace(output), 500),
			ToolStatus:   strings.TrimSpace(status),
			ToolExitCode: &code,
			ToolSuccess:  &success,
		}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}

	case "function_call":
		name, _ := item["name"].(string)
		status, _ := item["status"].(string)
		output, _ := item["output"].(string)
		success := codexToolSuccess(status, nil)
		slog.Debug("codexSession: function_call completed",
			"name", name, "status", status, "output_len", len(output),
		)
		evt := core.Event{
			Type:        core.EventToolResult,
			ToolName:    name,
			ToolResult:  truncate(strings.TrimSpace(output), 500),
			ToolStatus:  strings.TrimSpace(status),
			ToolSuccess: &success,
		}
		select {
		case cs.events <- evt:
		case <-cs.ctx.Done():
			return
		}

	case "function_call_output":
		slog.Debug("codexSession: function_call_output")

	case "error":
		msg, _ := item["message"].(string)
		if msg != "" && !strings.Contains(msg, "Falling back") {
			slog.Warn("codexSession: item error", "message", msg)
		}

	default:
		if toolName, known := codexToolNames[itemType]; known {
			input := codexExtractToolInput(item)
			evt := core.Event{Type: core.EventToolUse, ToolName: toolName, ToolInput: input}
			select {
			case cs.events <- evt:
			case <-cs.ctx.Done():
				return
			}
		} else {
			slog.Debug("codexSession: unhandled item type", "item_type", itemType)
		}
	}
}

// codexExtractToolInput extracts a human-readable input from a Codex tool item.
// For web_search, it reads action.queries[] or falls back to the top-level query.
func codexExtractToolInput(item map[string]any) string {
	if action, ok := item["action"].(map[string]any); ok {
		if queries, ok := action["queries"].([]any); ok && len(queries) > 0 {
			var parts []string
			for _, q := range queries {
				if s, ok := q.(string); ok && s != "" {
					parts = append(parts, s)
				}
			}
			if len(parts) > 0 {
				return strings.Join(parts, "\n")
			}
		}
		if q, _ := action["query"].(string); q != "" {
			return q
		}
	}
	if q, _ := item["query"].(string); q != "" {
		return q
	}
	if n, _ := item["name"].(string); n != "" {
		return n
	}
	return ""
}

func codexToolSuccess(status string, exitCode *int) bool {
	s := strings.ToLower(strings.TrimSpace(status))
	if exitCode != nil {
		return *exitCode == 0
	}
	return s == "completed" || s == "success" || s == "succeeded" || s == "ok"
}

func loadCodexRuntimeConfig(ctx context.Context, workDir string, extraEnv []string) (string, string, error) {
	cmd := exec.CommandContext(ctx, "codex", "app-server")
	cmd.Dir = workDir
	prepareCmdForKill(cmd)
	if len(extraEnv) > 0 {
		cmd.Env = core.MergeEnv(os.Environ(), extraEnv)
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return "", "", fmt.Errorf("runtime config stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", "", fmt.Errorf("runtime config stdout pipe: %w", err)
	}
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	if err := cmd.Start(); err != nil {
		return "", "", fmt.Errorf("runtime config start app-server: %w", err)
	}
	defer func() {
		_ = stdin.Close()
		if cmd.Process != nil {
			_ = cmd.Process.Kill()
		}
		_ = cmd.Wait()
	}()

	reader := bufio.NewReader(stdout)
	nextID := int64(1)

	if err := rpcRequestOverIO(stdin, reader, nextID, "initialize", map[string]any{
		"clientInfo": map[string]any{
			"name":    "cc-connect-codex-runtime-config",
			"title":   "CC Connect Codex Runtime Config",
			"version": "0.1.0",
		},
	}, nil); err != nil {
		return "", "", err
	}
	nextID++

	if err := rpcNotifyOverIO(stdin, "initialized", map[string]any{}); err != nil {
		return "", "", err
	}

	var resp struct {
		Config struct {
			Model                string  `json:"model"`
			ModelReasoningEffort *string `json:"model_reasoning_effort"`
		} `json:"config"`
	}
	if err := rpcRequestOverIO(stdin, reader, nextID, "config/read", map[string]any{
		"includeLayers": false,
	}, &resp); err != nil {
		return "", "", err
	}

	return strings.TrimSpace(resp.Config.Model), normalizeRuntimeReasoningEffort(stringValue(resp.Config.ModelReasoningEffort)), nil
}

func rpcRequestOverIO(stdin io.Writer, reader *bufio.Reader, id int64, method string, params any, out any) error {
	payload := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"method":  method,
		"params":  params,
	}
	if err := writeRPCMessage(stdin, payload); err != nil {
		return err
	}

	for {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			return fmt.Errorf("%s read response: %w", method, err)
		}

		var probe map[string]json.RawMessage
		if err := json.Unmarshal(bytes.TrimSpace(line), &probe); err != nil {
			continue
		}
		if _, ok := probe["id"]; !ok {
			continue
		}

		var resp rpcResponseEnvelope
		if err := json.Unmarshal(bytes.TrimSpace(line), &resp); err != nil {
			continue
		}
		respID, ok := rpcIDToInt64(resp.ID)
		if !ok || respID != id {
			continue
		}
		if resp.Error != nil {
			return fmt.Errorf("%s: %s", method, strings.TrimSpace(resp.Error.Message))
		}
		if out != nil {
			if err := json.Unmarshal(resp.Result, out); err != nil {
				return fmt.Errorf("%s decode response: %w", method, err)
			}
		}
		return nil
	}
}

func rpcNotifyOverIO(stdin io.Writer, method string, params any) error {
	payload := map[string]any{
		"jsonrpc": "2.0",
		"method":  method,
		"params":  params,
	}
	return writeRPCMessage(stdin, payload)
}

func writeRPCMessage(w io.Writer, payload any) error {
	b, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("encode rpc message: %w", err)
	}
	if _, err := w.Write(append(b, '\n')); err != nil {
		return fmt.Errorf("write rpc message: %w", err)
	}
	return nil
}

// RespondPermission is a no-op for Codex — permissions are handled via CLI flags.
func (cs *codexSession) RespondPermission(_ string, _ core.PermissionResult) error {
	return nil
}

func (cs *codexSession) Events() <-chan core.Event {
	return cs.events
}

func (cs *codexSession) CurrentSessionID() string {
	v, _ := cs.threadID.Load().(string)
	return v
}

func (cs *codexSession) GetWorkDir() string {
	return cs.workDir
}

func (cs *codexSession) GetModel() string {
	if model := strings.TrimSpace(cs.model); model != "" {
		return model
	}
	model, _ := cs.runtimeConfig()
	return model
}

func (cs *codexSession) GetReasoningEffort() string {
	if effort := strings.TrimSpace(cs.effort); effort != "" {
		return effort
	}
	_, effort := cs.runtimeConfig()
	return effort
}

func (cs *codexSession) Alive() bool {
	return cs.alive.Load()
}

func (cs *codexSession) GetContextUsage() *core.ContextUsage {
	cs.contextMu.RLock()
	defer cs.contextMu.RUnlock()
	return cloneContextUsage(cs.contextUsage)
}

func (cs *codexSession) runtimeConfig() (string, string) {
	cs.runtimeCfgMu.Lock()
	defer cs.runtimeCfgMu.Unlock()

	if !cs.runtimeCfgFetched.IsZero() && time.Since(cs.runtimeCfgFetched) < codexRuntimeConfigCacheTTL {
		return cs.runtimeCfgModel, cs.runtimeCfgEffort
	}

	ctx, cancel := context.WithTimeout(cs.ctx, codexRuntimeConfigTimeout)
	defer cancel()

	model, effort, err := loadCodexRuntimeConfig(ctx, cs.workDir, cs.extraEnv)
	if err == nil {
		cs.runtimeCfgModel = model
		cs.runtimeCfgEffort = effort
		cs.runtimeCfgFetchErr = nil
		cs.runtimeCfgFetched = time.Now()
		return model, effort
	}

	cs.runtimeCfgFetchErr = err
	if !cs.runtimeCfgFetched.IsZero() {
		return cs.runtimeCfgModel, cs.runtimeCfgEffort
	}
	return "", ""
}

func (cs *codexSession) refreshContextUsageFromRollout() {
	sessionID := strings.TrimSpace(cs.CurrentSessionID())
	if sessionID == "" {
		return
	}

	for attempt := 0; attempt < codexContextUsageRetryCount; attempt++ {
		cs.contextMu.RLock()
		cachedPath := cs.sessionFile
		cs.contextMu.RUnlock()

		usage, path, err := loadContextUsageFromRollout(cs.extraEnv, sessionID, cachedPath)
		if err == nil && usage != nil {
			cs.contextMu.Lock()
			cs.sessionFile = path
			cs.contextUsage = cloneContextUsage(usage)
			cs.contextMu.Unlock()
			return
		}

		if attempt == codexContextUsageRetryCount-1 {
			if err != nil {
				slog.Debug("codexSession: context usage unavailable", "thread_id", sessionID, "error", err)
			}
			return
		}

		select {
		case <-time.After(codexContextUsageRetryDelay):
		case <-cs.ctx.Done():
			return
		}
	}
}

func (cs *codexSession) Close() error {
	cs.alive.Store(false)
	cs.cancel()
	done := make(chan struct{})
	go func() {
		cs.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
		// readLoop has exited; safe to close the events channel.
		cs.closeOnce.Do(func() {
			close(cs.events)
		})
		return nil
	case <-time.After(codexSessionCloseTimeout):
		cmds := cs.activeCmds()
		slog.Warn("codexSession: graceful close timed out, killing active process groups",
			"wait", codexSessionCloseTimeout,
			"count", len(cmds))
		if err := forceKillAllCmds(cmds); err != nil {
			slog.Debug("codexSession: force kill failed", "error", err)
		}
		select {
		case <-done:
			cs.closeOnce.Do(func() {
				close(cs.events)
			})
			return nil
		case <-time.After(codexSessionForceKillWait):
			// Do not close(cs.events) here: readLoop may still be in handleEvent
			// (e.g. turn.completed -> flushPendingAsText) and would panic on send.
			slog.Warn("codexSession: force kill wait timed out, deferring events channel close until readLoop exits",
				"wait", codexSessionForceKillWait)
			go func() {
				<-done
				cs.closeOnce.Do(func() {
					close(cs.events)
				})
			}()
			return nil
		}
	}
}

func (cs *codexSession) addCmd(cmd *exec.Cmd) {
	cs.cmdMu.Lock()
	defer cs.cmdMu.Unlock()
	cs.cmds[cmd] = struct{}{}
}

func (cs *codexSession) removeCmd(cmd *exec.Cmd) {
	cs.cmdMu.Lock()
	defer cs.cmdMu.Unlock()
	delete(cs.cmds, cmd)
}

func (cs *codexSession) activeCmds() []*exec.Cmd {
	cs.cmdMu.Lock()
	defer cs.cmdMu.Unlock()
	cmds := make([]*exec.Cmd, 0, len(cs.cmds))
	for cmd := range cs.cmds {
		cmds = append(cmds, cmd)
	}
	return cmds
}

func forceKillAllCmds(cmds []*exec.Cmd) error {
	var errs []error
	for _, cmd := range cmds {
		if err := forceKillCmd(cmd); err != nil {
			errs = append(errs, err)
		}
	}
	return errors.Join(errs...)
}

// extractItemText extracts text from an item's array field (e.g. "summary" or "content").
// It looks for elements matching the given elementType and concatenates their "text" fields.
// Falls back to the item's top-level "text" field if the array is missing or empty.
func extractItemText(item map[string]any, arrayField, elementType string) string {
	if arr, ok := item[arrayField].([]any); ok {
		var parts []string
		for _, elem := range arr {
			m, ok := elem.(map[string]any)
			if !ok {
				continue
			}
			if elementType != "" {
				if t, _ := m["type"].(string); t != elementType {
					continue
				}
			}
			if t, _ := m["text"].(string); t != "" {
				parts = append(parts, t)
			}
		}
		if len(parts) > 0 {
			return strings.Join(parts, "\n")
		}
	}
	text, _ := item["text"].(string)
	return text
}

func truncate(s string, maxRunes int) string {
	if utf8.RuneCountInString(s) <= maxRunes {
		return s
	}
	return string([]rune(s)[:maxRunes]) + "..."
}