iop/apps/node/internal/adapters/cli/codex_app_server_process.go

package cli

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"sync"
	"sync/atomic"

	"iop/packages/go/config"
)

type codexAppServerProc struct {
	cmd       *exec.Cmd
	stdin     io.WriteCloser
	stdout    *bufio.Scanner
	notifCh   chan appServerNotification
	done      chan struct{}
	closeOnce sync.Once
	nextID    atomic.Int64
	pending   sync.Map // ID -> chan appServerResponse
}

type appServerRequest struct {
	ID     int64          `json:"id"`
	Method string         `json:"method"`
	Params map[string]any `json:"params"`
}

type appServerResponse struct {
	ID     int64          `json:"id"`
	Result map[string]any `json:"result,omitempty"`
	Error  *appServerErr  `json:"error,omitempty"`
}

type appServerErr struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type appServerNotification struct {
	Method string         `json:"method"`
	Params map[string]any `json:"params,omitempty"`
}

func startCodexAppServerProc(ctx context.Context, profile config.CLIProfileConf, workspace string) (*codexAppServerProc, error) {
	args := append(append([]string{}, profile.Args...), "--stdio")
	cmd, err := buildCmdWithoutContext(profile.Command, args, profile.Env, workspace)
	if err != nil {
		return nil, err
	}

	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdin pipe: %w", err)
	}
	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}
	// stderr is intentionally discarded; process exit errors surface via done channel.
	cmd.Stderr = nil

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

	p := &codexAppServerProc{
		cmd:     cmd,
		stdin:   stdinPipe,
		stdout:  bufio.NewScanner(stdoutPipe),
		notifCh: make(chan appServerNotification, 32),
		done:    make(chan struct{}),
	}
	p.stdout.Buffer(make([]byte, 64*1024), 4*1024*1024)

	go p.readLoop()
	go func() {
		_ = cmd.Wait()
		p.closeOnce.Do(func() { close(p.done) })
	}()

	return p, nil
}

// readLoop dispatches incoming newline-delimited JSON to pending response channels
// or the notification channel.
func (p *codexAppServerProc) readLoop() {
	for p.stdout.Scan() {
		line := p.stdout.Bytes()
		if len(line) == 0 {
			continue
		}

		// Try response first (has "id" field as number).
		var resp appServerResponse
		if err := json.Unmarshal(line, &resp); err == nil && resp.ID != 0 {
			if ch, ok := p.pending.Load(resp.ID); ok {
				select {
				case ch.(chan appServerResponse) <- resp:
				default:
				}
			}
			continue
		}

		// Treat as notification. Blocking send ensures no notification is dropped;
		// the drain goroutine or a closed done channel will unblock this if needed.
		var notif appServerNotification
		if err := json.Unmarshal(line, &notif); err == nil && notif.Method != "" {
			select {
			case p.notifCh <- notif:
			case <-p.done:
				return
			}
		}
	}
	p.closeOnce.Do(func() { close(p.done) })
}

// recvResponse waits for a response on ch, giving priority to an already-arrived
// response over p.done or ctx cancellation. This prevents the case where the
// process sends its response and then immediately exits, and Go's random select
// picks p.done instead of ch.
func (p *codexAppServerProc) recvResponse(ctx context.Context, ch chan appServerResponse) (appServerResponse, error) {
	// Fast path: response already in buffer.
	select {
	case resp := <-ch:
		return resp, nil
	default:
	}
	// Slow path: wait for any of the three signals.
	select {
	case resp := <-ch:
		return resp, nil
	case <-p.done:
		// One last drain: response may have arrived between the default check and here.
		select {
		case resp := <-ch:
			return resp, nil
		default:
		}
		return appServerResponse{}, fmt.Errorf("codex app-server process exited")
	case <-ctx.Done():
		return appServerResponse{}, ctx.Err()
	}
}

func (p *codexAppServerProc) send(ctx context.Context, method string, params map[string]any) (appServerResponse, error) {
	id := p.nextID.Add(1)
	if params == nil {
		params = map[string]any{}
	}
	req := appServerRequest{ID: id, Method: method, Params: params}

	raw, err := json.Marshal(req)
	if err != nil {
		return appServerResponse{}, fmt.Errorf("marshal request: %w", err)
	}
	raw = append(raw, '\n')

	ch := make(chan appServerResponse, 1)
	p.pending.Store(id, ch)
	defer p.pending.Delete(id)

	// Write stdin in a goroutine so that a blocked pipe write does not prevent
	// us from observing process exit or an early response via p.done / ctx.
	writeErr := make(chan error, 1)
	go func() {
		_, err := p.stdin.Write(raw)
		writeErr <- err
	}()

	// Wait for write completion, but also watch for an early response or exit.
	// A fast server may respond and exit before the write goroutine reports back.
	select {
	case resp, ok := <-ch:
		// Response arrived before write confirmed — still valid.
		if ok {
			if resp.Error != nil {
				return resp, fmt.Errorf("rpc error %d: %s", resp.Error.Code, resp.Error.Message)
			}
			return resp, nil
		}
	default:
	}
	select {
	case err := <-writeErr:
		if err != nil {
			return appServerResponse{}, fmt.Errorf("write stdin: %w", err)
		}
	case resp := <-ch:
		if resp.Error != nil {
			return resp, fmt.Errorf("rpc error %d: %s", resp.Error.Code, resp.Error.Message)
		}
		return resp, nil
	case <-p.done:
		// Check for response that arrived simultaneously with exit.
		select {
		case resp := <-ch:
			if resp.Error != nil {
				return resp, fmt.Errorf("rpc error %d: %s", resp.Error.Code, resp.Error.Message)
			}
			return resp, nil
		default:
		}
		return appServerResponse{}, fmt.Errorf("codex app-server process exited")
	case <-ctx.Done():
		return appServerResponse{}, ctx.Err()
	}

	resp, err := p.recvResponse(ctx, ch)
	if err != nil {
		return appServerResponse{}, err
	}
	if resp.Error != nil {
		return resp, fmt.Errorf("rpc error %d: %s", resp.Error.Code, resp.Error.Message)
	}
	return resp, nil
}

// sendNotification sends a JSON-RPC notification (no id, no response expected).
func (p *codexAppServerProc) sendNotification(method string, params map[string]any) error {
	type notifMsg struct {
		Method string         `json:"method"`
		Params map[string]any `json:"params,omitempty"`
	}
	raw, err := json.Marshal(notifMsg{Method: method, Params: params})
	if err != nil {
		return fmt.Errorf("marshal notification: %w", err)
	}
	raw = append(raw, '\n')
	_, err = p.stdin.Write(raw)
	return err
}

func (p *codexAppServerProc) close() {
	_ = p.stdin.Close()
	if p.cmd != nil && p.cmd.Process != nil {
		_ = p.cmd.Process.Kill()
	}
	p.closeOnce.Do(func() { close(p.done) })
}

// codexAppServerInit runs initialize → initialized → thread/start and returns the thread id.
func codexAppServerInit(ctx context.Context, p *codexAppServerProc) (string, error) {
	initParams := map[string]any{
		"protocolVersion": "2024-11-05",
		"clientInfo":      map[string]any{"name": "iop-node", "version": "0"},
		"capabilities":    map[string]any{},
	}
	if _, err := p.send(ctx, "initialize", initParams); err != nil {
		return "", fmt.Errorf("initialize: %w", err)
	}

	if err := p.sendNotification("initialized", nil); err != nil {
		return "", fmt.Errorf("initialized notification: %w", err)
	}

	resp, err := p.send(ctx, "thread/start", map[string]any{})
	if err != nil {
		return "", fmt.Errorf("thread/start: %w", err)
	}

	// Actual shape: result.thread.id (ThreadStartResponse)
	threadID := nestedString(resp.Result, "thread", "id")
	// Legacy/flat fallbacks for compatibility with older or fake servers.
	if threadID == "" {
		threadID, _ = resp.Result["threadId"].(string)
	}
	if threadID == "" {
		threadID, _ = resp.Result["id"].(string)
	}
	if threadID == "" {
		return "", fmt.Errorf("thread/start: server did not return a thread id")
	}
	return threadID, nil
}