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

package cli

import (
	"context"
	"fmt"
	"sync"
	"time"

	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

// codexAppServerSession tracks a single long-lived `codex app-server --stdio` process
// bound to a (target, sessionID) key. The process is started lazily on first use and
// reused across Execute calls until it exits or is terminated.
type codexAppServerSession struct {
	key      sessionKey
	threadID string // Codex thread id assigned by thread/start

	mu     sync.Mutex
	proc   *codexAppServerProc
	closed bool
}

// codexAppServerExecutor is the sessionReporter + executor for modeCodexAppServer profiles.
type codexAppServerExecutor struct {
	cli      *CLI
	mu       sync.Mutex
	sessions map[sessionKey]*codexAppServerSession
}

type codexAppServerRun struct {
	executor *codexAppServerExecutor
	ctx      context.Context
	runID    string
	threadID string
	turnID   string
	proc     *codexAppServerProc
	sink     runtime.EventSink
}

func (r *codexAppServerRun) drain() error {
	for {
		select {
		case <-r.ctx.Done():
			return r.ctx.Err()
		case notif, ok := <-r.proc.notifCh:
			if !ok {
				return fmt.Errorf("codex app-server notification channel closed")
			}
			done, err := r.handleNotification(notif)
			if err != nil {
				return err
			}
			if done {
				return nil
			}
		case <-r.proc.done:
			return r.handleProcessExit()
		}
	}
}

func (r *codexAppServerRun) handleNotification(notif appServerNotification) (bool, error) {
	events, done, err := decodeAppServerNotification(notif, r.runID, r.threadID, r.turnID)
	for _, ev := range events {
		if emitErr := r.sink.Emit(r.ctx, ev); emitErr != nil {
			return false, emitErr
		}
	}
	return done, err
}

func (r *codexAppServerRun) handleProcessExit() error {
	// Process exited. Drain any notifications already in the buffer before
	// declaring an error — the terminal event may have arrived with the exit.
	for {
		select {
		case notif, ok := <-r.proc.notifCh:
			if !ok {
				return fmt.Errorf("codex app-server process exited during turn")
			}
			done, err := r.handleNotification(notif)
			if err != nil {
				return err
			}
			if done {
				return nil
			}
		default:
			return fmt.Errorf("codex app-server process exited during turn")
		}
	}
}

func (e *codexAppServerExecutor) Execute(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf, sink runtime.EventSink) error {
	sess, err := e.resolveCodexAppServerSession(ctx, spec, profile)
	if err != nil {
		return emitReturnedError(ctx, sink, spec.RunID, err)
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()

	if err := sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     spec.RunID,
		Type:      runtime.EventTypeStart,
		Timestamp: time.Now(),
	}); err != nil {
		return err
	}

	prompt := extractPrompt(spec.Input)
	turnID, err := codexAppServerTurnStart(ctx, sess.proc, sess.threadID, prompt)
	if err != nil {
		return err
	}

	run := &codexAppServerRun{
		executor: e,
		ctx:      ctx,
		runID:    spec.RunID,
		threadID: sess.threadID,
		turnID:   turnID,
		proc:     sess.proc,
		sink:     sink,
	}
	return run.drain()
}

// codexAppServerTurnStart sends turn/start with the prompt and returns the turn id.
func codexAppServerTurnStart(ctx context.Context, p *codexAppServerProc, threadID, prompt string) (string, error) {
	// TurnStartParams.input is required; UserInput text variant requires text_elements.
	params := map[string]any{
		"threadId": threadID,
		"input": []map[string]any{
			{"type": "text", "text": prompt, "text_elements": []any{}},
		},
	}
	resp, err := p.send(ctx, "turn/start", params)
	if err != nil {
		return "", fmt.Errorf("turn/start: %w", err)
	}
	// Actual shape: result.turn.id (TurnStartResponse)
	turnID := nestedString(resp.Result, "turn", "id")
	// Legacy/flat fallbacks.
	if turnID == "" {
		turnID, _ = resp.Result["turnId"].(string)
	}
	if turnID == "" {
		turnID, _ = resp.Result["id"].(string)
	}
	if turnID == "" {
		return "", fmt.Errorf("turn/start: server did not return a turn id")
	}
	return turnID, nil
}

func codexAppServerDrainNotifications(ctx context.Context, p *codexAppServerProc, runID, threadID, turnID string, sink runtime.EventSink) error {
	run := &codexAppServerRun{
		ctx:      ctx,
		runID:    runID,
		threadID: threadID,
		turnID:   turnID,
		proc:     p,
		sink:     sink,
	}
	return run.drain()
}

// decodeAppServerNotification maps a single app-server notification to RuntimeEvents.
// threadID and turnID are the execution-context values used as fallbacks when the
// notification params do not carry them. Returns (events, turnDone, error).
func decodeAppServerNotification(notif appServerNotification, runID, threadID, turnID string) ([]runtime.RuntimeEvent, bool, error) {
	// Prefer IDs from the notification params; fall back to execution-context values.
	effectiveThread := notif.Params["threadId"]
	if s, ok := effectiveThread.(string); ok && s != "" {
		threadID = s
	}
	effectiveTurn := notif.Params["turnId"]
	if s, ok := effectiveTurn.(string); ok && s != "" {
		turnID = s
	}

	meta := map[string]string{"source": "codex-app-server"}
	if threadID != "" {
		meta["thread_id"] = threadID
	}
	if turnID != "" {
		meta["turn_id"] = turnID
	}

	switch notif.Method {
	case "item/agentMessage/delta":
		// AgentMessageDeltaNotification: { threadId, turnId, itemId, delta }
		delta := nestedString(notif.Params, "delta")
		if delta == "" {
			return nil, false, nil
		}
		return []runtime.RuntimeEvent{{
			RunID:     runID,
			Type:      runtime.EventTypeDelta,
			Delta:     delta,
			Metadata:  meta,
			Timestamp: time.Now(),
		}}, false, nil

	case "item/completed":
		// Suppress duplicate final delta — already streamed via delta notifications.
		return nil, false, nil

	case "turn/completed":
		// TurnCompletedNotification: { threadId, turn: { id, status, error, ... } }
		// status may be "completed" | "interrupted" | "failed" | "inProgress"
		status := nestedString(notif.Params, "turn", "status")
		if status == "failed" {
			msg := nestedString(notif.Params, "turn", "error", "message")
			if msg == "" {
				msg = "turn failed"
			}
			return []runtime.RuntimeEvent{{
				RunID:     runID,
				Type:      runtime.EventTypeError,
				Error:     msg,
				Metadata:  meta,
				Timestamp: time.Now(),
			}}, false, fmt.Errorf("codex app-server turn failed: %s", msg)
		}
		return []runtime.RuntimeEvent{{
			RunID:     runID,
			Type:      runtime.EventTypeComplete,
			Metadata:  meta,
			Timestamp: time.Now(),
		}}, true, nil

	case "error":
		// ErrorNotification: { error: { message, ... }, willRetry, threadId, turnId }
		msg := nestedString(notif.Params, "error", "message")
		if msg == "" {
			// flat fallback for non-standard shapes
			msg = nestedString(notif.Params, "message")
		}
		if msg == "" {
			msg = "codex app-server error"
		}
		return []runtime.RuntimeEvent{{
			RunID:     runID,
			Type:      runtime.EventTypeError,
			Error:     msg,
			Metadata:  meta,
			Timestamp: time.Now(),
		}}, false, fmt.Errorf("codex app-server: %s", msg)

	case "turn/failed":
		// Legacy method name kept for compatibility.
		msg := nestedString(notif.Params, "error", "message")
		if msg == "" {
			msg = nestedString(notif.Params, "message")
		}
		if msg == "" {
			msg = "turn failed"
		}
		return []runtime.RuntimeEvent{{
			RunID:     runID,
			Type:      runtime.EventTypeError,
			Error:     msg,
			Metadata:  meta,
			Timestamp: time.Now(),
		}}, false, fmt.Errorf("codex app-server turn failed: %s", msg)
	}
	return nil, false, nil
}

// nestedString walks a params map by keys and returns the string value at the path.
func nestedString(params map[string]any, keys ...string) string {
	var cur any = params
	for i, k := range keys {
		m, ok := cur.(map[string]any)
		if !ok {
			return ""
		}
		cur = m[k]
		if i == len(keys)-1 {
			s, _ := cur.(string)
			return s
		}
	}
	return ""
}

func (e *codexAppServerExecutor) resolveCodexAppServerSession(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf) (*codexAppServerSession, error) {
	target := cliTargetName(spec)
	key := newSessionKey(spec)

	e.mu.Lock()
	if sess, ok := e.sessions[key]; ok {
		e.mu.Unlock()
		return sess, nil
	}
	if spec.SessionMode == runtime.SessionModeRequireExisting {
		e.mu.Unlock()
		return nil, fmt.Errorf("cli adapter: no codex app-server session for target %q session %q", target, key.sessionID)
	}
	e.mu.Unlock()

	proc, err := startCodexAppServerProc(ctx, profile, spec.Workspace)
	if err != nil {
		return nil, fmt.Errorf("cli adapter: start codex app-server: %w", err)
	}

	threadID, err := codexAppServerInit(ctx, proc)
	if err != nil {
		proc.close()
		return nil, fmt.Errorf("cli adapter: codex app-server init: %w", err)
	}

	sess := &codexAppServerSession{
		key:      key,
		threadID: threadID,
		proc:     proc,
	}

	e.mu.Lock()
	if existing, ok := e.sessions[key]; ok {
		e.mu.Unlock()
		proc.close()
		return existing, nil
	}
	e.sessions[key] = sess
	e.mu.Unlock()
	return sess, nil
}

func (e *codexAppServerExecutor) Sessions() []sessionListEntry {
	e.mu.Lock()
	defer e.mu.Unlock()
	snaps := make([]sessionListEntry, 0, len(e.sessions))
	for k := range e.sessions {
		snaps = append(snaps, sessionListEntry{modeCodexAppServer, k.target, k.sessionID, k.workspace})
	}
	return snaps
}

// Terminate closes every workspace variant of the given target/sessionID.
func (e *codexAppServerExecutor) Terminate(_ context.Context, target, sessionID string) (bool, error) {
	sid := normalizeSessionID(sessionID)
	e.mu.Lock()
	var matched []*codexAppServerSession
	for k, sess := range e.sessions {
		if k.target == target && k.sessionID == sid {
			matched = append(matched, sess)
			delete(e.sessions, k)
		}
	}
	e.mu.Unlock()
	if len(matched) == 0 {
		return false, nil
	}
	for _, sess := range matched {
		closeCodexAppServerSession(sess)
	}
	return true, nil
}

func (e *codexAppServerExecutor) Stop(_ context.Context) error {
	e.mu.Lock()
	cp := make(map[sessionKey]*codexAppServerSession, len(e.sessions))
	for k, v := range e.sessions {
		cp[k] = v
	}
	e.sessions = make(map[sessionKey]*codexAppServerSession)
	e.mu.Unlock()
	for _, sess := range cp {
		closeCodexAppServerSession(sess)
	}
	return nil
}

func closeCodexAppServerSession(sess *codexAppServerSession) {
	if sess == nil {
		return
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()
	if sess.closed {
		return
	}
	sess.closed = true
	if sess.proc != nil {
		sess.proc.close()
	}
}