yao/agent/caller/orchestrator.go

package caller

import (
	"fmt"
	"sync"

	agentContext "github.com/yaoapp/yao/agent/context"
	"github.com/yaoapp/yao/trace/types"
)

// Orchestrator handles parallel agent calls with different concurrency patterns
// Modeled after JavaScript Promise patterns (all, any, race)
type Orchestrator struct {
	ctx *agentContext.Context
}

// NewOrchestrator creates a new Orchestrator for parallel agent calls
func NewOrchestrator(ctx *agentContext.Context) *Orchestrator {
	return &Orchestrator{ctx: ctx}
}

// callResult is used internally to pass results through channels
type callResult struct {
	idx    int
	result *Result
}

// All executes all agent calls and waits for all to complete (like Promise.all)
// Returns results in the same order as requests, regardless of completion order
// Each call uses a forked context to avoid race conditions on shared state
func (o *Orchestrator) All(reqs []*Request) []*Result {
	if len(reqs) == 0 {
		return []*Result{}
	}

	results := make([]*Result, len(reqs))
	var wg sync.WaitGroup
	var mu sync.Mutex

	for i, req := range reqs {
		wg.Add(1)
		go func(idx int, r *Request) {
			defer wg.Done()
			defer func() {
				if err := recover(); err != nil {
					mu.Lock()
					results[idx] = &Result{
						AgentID: r.AgentID,
						Error:   "agent call panic recovered",
					}
					mu.Unlock()
				}
			}()

			// Use forked context to avoid race conditions
			result := o.callAgentWithForkedContext(r)
			mu.Lock()
			results[idx] = result
			mu.Unlock()
		}(i, req)
	}

	wg.Wait()
	return results
}

// Any returns as soon as any agent call succeeds (has non-error result) (like Promise.any)
// Other calls continue in background but results are discarded after first success
// Returns all results received so far when first success is found
// Each call uses a forked context to avoid race conditions on shared state
func (o *Orchestrator) Any(reqs []*Request) []*Result {
	if len(reqs) == 0 {
		return []*Result{}
	}

	results := make([]*Result, len(reqs))
	resultChan := make(chan callResult, len(reqs))

	var wg sync.WaitGroup
	done := make(chan struct{})

	for i, req := range reqs {
		wg.Add(1)
		go func(idx int, r *Request) {
			defer wg.Done()
			defer func() {
				if err := recover(); err != nil {
					// Send panic result through channel
					select {
					case <-done:
					case resultChan <- callResult{idx: idx, result: &Result{
						AgentID: r.AgentID,
						Error:   "agent call panic recovered",
					}}:
					}
				}
			}()

			// Check if done before starting
			select {
			case <-done:
				return
			default:
			}

			// Use forked context to avoid race conditions
			result := o.callAgentWithForkedContext(r)

			// Try to send result
			select {
			case <-done:
				// Already found a successful result
			case resultChan <- callResult{idx: idx, result: result}:
			}
		}(i, req)
	}

	// Close channel when all goroutines complete
	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Collect results until we find one with success (no error and has content)
	var foundSuccess bool
	for res := range resultChan {
		results[res.idx] = res.result
		// Check if this result is successful (no error)
		if !foundSuccess && res.result != nil && res.result.Error == "" {
			foundSuccess = true
			close(done) // Signal other goroutines to stop
		}
	}

	return results
}

// Race returns as soon as any agent call completes (like Promise.race)
// Returns immediately when first result arrives, regardless of success/failure
// Note: Still waits for all goroutines to complete before returning to avoid resource leaks
// Each call uses a forked context to avoid race conditions on shared state
func (o *Orchestrator) Race(reqs []*Request) []*Result {
	if len(reqs) == 0 {
		return []*Result{}
	}

	results := make([]*Result, len(reqs))
	resultChan := make(chan callResult, len(reqs))

	var wg sync.WaitGroup
	done := make(chan struct{})

	for i, req := range reqs {
		wg.Add(1)
		go func(idx int, r *Request) {
			defer wg.Done()
			defer func() {
				if err := recover(); err != nil {
					// Send panic result through channel
					select {
					case <-done:
					case resultChan <- callResult{idx: idx, result: &Result{
						AgentID: r.AgentID,
						Error:   "agent call panic recovered",
					}}:
					}
				}
			}()

			// Check if done before starting
			select {
			case <-done:
				return
			default:
			}

			// Use forked context to avoid race conditions
			result := o.callAgentWithForkedContext(r)

			// Try to send result
			select {
			case <-done:
				// Already got first result
			case resultChan <- callResult{idx: idx, result: result}:
			}
		}(i, req)
	}

	// Close channel when all goroutines complete
	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Get first result and signal others to stop
	var gotFirst bool
	for res := range resultChan {
		results[res.idx] = res.result
		if !gotFirst {
			gotFirst = true
			close(done) // Signal other goroutines to stop
		}
	}

	return results
}

// callAgent executes a single agent call using the AgentGetterFunc
// This method handles context sharing and result extraction
func (o *Orchestrator) callAgent(req *Request) *Result {
	return o.callAgentWithContext(o.ctx, req)
}

// callAgentWithForkedContext executes a single agent call with a forked context
// This is used by batch operations (All/Any/Race) to avoid race conditions
// when multiple goroutines modify shared context state (Stack, Logger, etc.)
func (o *Orchestrator) callAgentWithForkedContext(req *Request) *Result {
	// Fork the context to get independent Stack and Logger
	forkedCtx := o.ctx.Fork()
	return o.callAgentWithContext(forkedCtx, req)
}

// callAgentWithContext executes a single agent call with the given context
// This is the core implementation used by both callAgent and callAgentWithForkedContext
func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Request) *Result {
	if req == nil {
		return &Result{Error: "nil request"}
	}

	// Get the agent using the getter function
	if AgentGetterFunc == nil {
		return NewResult(req.AgentID, nil, fmt.Errorf("agent getter not initialized"))
	}

	agent, err := AgentGetterFunc(req.AgentID)
	if err != nil {
		return NewResult(req.AgentID, nil, fmt.Errorf("failed to get agent: %w", err))
	}

	// Mark this as an agent-to-agent fork call for proper source tracking
	// RefererAgentFork distinguishes ctx.agent.Call from delegate calls
	ctx.Referer = agentContext.RefererAgentFork

	// Build context options for the call
	var ctxOpts *agentContext.Options
	if req.Options != nil {
		ctxOpts = req.Options.ToContextOptions()
	} else {
		ctxOpts = &agentContext.Options{}
	}

	// If request has a handler, set OnMessage callback
	if req.Handler != nil {
		if ctxOpts == nil {
			ctxOpts = &agentContext.Options{}
		}
		// Set OnMessage to receive SSE messages
		ctxOpts.OnMessage = req.Handler
	}

	// Add trace node for A2A call using the ORIGINAL parent context's trace
	// (forked contexts have nil trace and nil Stack, so ctx.Trace() would create a new orphan trace)
	parentTrace, _ := o.ctx.Trace()
	var a2aNode types.Node
	if parentTrace != nil {
		a2aNode, _ = parentTrace.Add(
			map[string]any{
				"agent_id": req.AgentID,
				"referer":  string(ctx.Referer),
			},
			types.TraceNodeOption{
				Label:       fmt.Sprintf("Agent: %s", req.AgentID),
				Type:        "agent_call",
				Icon:        "smart_toy",
				Description: fmt.Sprintf("A2A call to '%s'", req.AgentID),
			},
		)
	}

	resp, err := agent.Stream(ctx, req.Messages, ctxOpts)
	if err != nil {
		if a2aNode != nil {
			a2aNode.Fail(err)
		}
		return NewResult(req.AgentID, nil, fmt.Errorf("agent call failed: %w", err))
	}

	if a2aNode != nil {
		a2aNode.Complete(map[string]any{
			"agent_id": req.AgentID,
			"status":   "completed",
		})
	}

	return NewResult(req.AgentID, resp, nil)
}

// extractContentFromCompletion extracts the text content from a completion response
func extractContentFromCompletion(completion *agentContext.CompletionResponse) string {
	if completion == nil {
		return ""
	}

	// Content can be string or []ContentPart
	switch content := completion.Content.(type) {
	case string:
		return content
	case []interface{}:
		// Handle array of content parts - extract text parts
		var texts []string
		for _, part := range content {
			if partMap, ok := part.(map[string]interface{}); ok {
				if partType, ok := partMap["type"].(string); ok && partType == "text" {
					if text, ok := partMap["text"].(string); ok {
						texts = append(texts, text)
					}
				}
			}
		}
		if len(texts) > 0 {
			return texts[0] // Return first text content
		}
	}

	return ""
}