thread.go

package codex

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
)

// Turn is a completed turn.
type Turn struct {
	// Items contains completed items emitted during the turn.
	Items []ThreadItem
	// FinalResponse is the latest completed agent message text, if any.
	FinalResponse string
	// Usage contains token usage from the completed turn.
	Usage *Usage
}

// RunResult is an alias for Turn.
type RunResult = Turn

// RunStreamedResult provides a stream of events and a terminal wait handle.
type RunStreamedResult struct {
	Events <-chan ThreadEvent
	done   <-chan error
}

// Wait waits for the streamed turn to finish and returns any terminal error.
func (r *RunStreamedResult) Wait() error {
	if r == nil || r.done == nil {
		return nil
	}
	return <-r.done
}

// UserInput represents a structured input segment for a turn.
type UserInput struct {
	// Type is either `text` or `local_image`.
	Type string `json:"type"`
	// Text contains prompt text for `text` items.
	Text string `json:"text,omitempty"`
	// Path contains the local image path for `local_image` items.
	Path string `json:"path,omitempty"`
}

// Input is either a prompt string or a slice of UserInput values.
type Input = any

// Thread represents a conversation with the agent.
type Thread struct {
	exec          *CodexExec
	options       CodexOptions
	threadOptions ThreadOptions

	mu sync.RWMutex
	id string
}

func newThread(exec *CodexExec, options CodexOptions, threadOptions ThreadOptions, id string) *Thread {
	return &Thread{
		exec:          exec,
		options:       options,
		threadOptions: threadOptions,
		id:            id,
	}
}

// ID returns the thread ID, populated after the first thread.started event.
func (t *Thread) ID() string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.id
}

func (t *Thread) setID(id string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.id = id
}

// RunStreamed starts a turn and streams structured events.
func (t *Thread) RunStreamed(input Input, options ...TurnOptions) (*RunStreamedResult, error) {
	var turnOptions TurnOptions
	if len(options) > 0 {
		turnOptions = options[0]
	}

	schemaPath, cleanup, err := createOutputSchemaFile(turnOptions.OutputSchema)
	if err != nil {
		return nil, err
	}

	prompt, images, err := normalizeInput(input)
	if err != nil {
		_ = cleanup()
		return nil, err
	}

	ctx := turnOptions.Context
	if ctx == nil {
		ctx = context.Background()
	}

	events := make(chan ThreadEvent)
	done := make(chan error, 1)

	go func() {
		defer close(events)
		defer close(done)
		defer func() {
			_ = cleanup()
		}()

		err := t.exec.Run(ctx, CodexExecArgs{
			Input:                 prompt,
			BaseURL:               t.options.BaseURL,
			APIKey:                t.options.APIKey,
			ThreadID:              t.ID(),
			Images:                images,
			Model:                 t.threadOptions.Model,
			SandboxMode:           t.threadOptions.SandboxMode,
			WorkingDirectory:      t.threadOptions.WorkingDirectory,
			SkipGitRepoCheck:      t.threadOptions.SkipGitRepoCheck,
			OutputSchemaFile:      schemaPath,
			ModelReasoningEffort:  t.threadOptions.ModelReasoningEffort,
			NetworkAccessEnabled:  t.threadOptions.NetworkAccessEnabled,
			WebSearchMode:         t.threadOptions.WebSearchMode,
			WebSearchEnabled:      t.threadOptions.WebSearchEnabled,
			ApprovalPolicy:        t.threadOptions.ApprovalPolicy,
			AdditionalDirectories: t.threadOptions.AdditionalDirectories,
		}, func(line string) error {
			var event ThreadEvent
			if err := json.Unmarshal([]byte(line), &event); err != nil {
				return fmt.Errorf("failed to parse item: %s: %w", line, err)
			}
			if event.Type == "thread.started" && event.ThreadID != "" {
				t.setID(event.ThreadID)
			}
			select {
			case events <- event:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			}
		})
		done <- err
	}()

	return &RunStreamedResult{Events: events, done: done}, nil
}

// Run starts a turn and buffers the completed result.
func (t *Thread) Run(input Input, options ...TurnOptions) (Turn, error) {
	stream, err := t.RunStreamed(input, options...)
	if err != nil {
		return Turn{}, err
	}

	items := make([]ThreadItem, 0)
	var finalResponse string
	var usage *Usage
	var turnFailure *ThreadError

	for event := range stream.Events {
		switch event.Type {
		case "item.completed":
			if event.Item != nil {
				if event.Item.Type == "agent_message" {
					finalResponse = event.Item.Text
				}
				items = append(items, *event.Item)
			}
		case "turn.completed":
			usage = event.Usage
		case "turn.failed":
			turnFailure = event.Error
		}
	}

	if err := stream.Wait(); err != nil {
		return Turn{}, err
	}
	if turnFailure != nil {
		return Turn{}, errors.New(turnFailure.Message)
	}

	return Turn{
		Items:         items,
		FinalResponse: finalResponse,
		Usage:         usage,
	}, nil
}

func normalizeInput(input Input) (string, []string, error) {
	switch value := input.(type) {
	case string:
		return value, nil, nil
	case []UserInput:
		return normalizeStructuredInput(value)
	case []*UserInput:
		items := make([]UserInput, 0, len(value))
		for _, item := range value {
			if item == nil {
				return "", nil, fmt.Errorf("input items cannot be nil")
			}
			items = append(items, *item)
		}
		return normalizeStructuredInput(items)
	default:
		return "", nil, fmt.Errorf("input must be a string or []UserInput")
	}
}

func normalizeStructuredInput(input []UserInput) (string, []string, error) {
	promptParts := make([]string, 0, len(input))
	images := make([]string, 0)
	for _, item := range input {
		switch item.Type {
		case "text":
			promptParts = append(promptParts, item.Text)
		case "local_image":
			images = append(images, item.Path)
		default:
			return "", nil, fmt.Errorf("unsupported input type: %s", item.Type)
		}
	}
	return joinPromptParts(promptParts), images, nil
}

func joinPromptParts(parts []string) string {
	if len(parts) == 0 {
		return ""
	}
	result := parts[0]
	for i := 1; i < len(parts); i++ {
		result += "\n\n" + parts[i]
	}
	return result
}