ragflow/internal/harness/core/workflow.go

package core

import (
	"context"
	"fmt"
	"runtime/debug"
	"sync"

	"ragflow/internal/harness/core/schema"
)

type workflowMode int

const (
	workflowModeUnknown  workflowMode = iota
	workflowModeSequential
	workflowModeLoop
	workflowModeParallel
)

type workflowState struct {
	InterruptIdx int
}
type workflowParallelState struct {
	SubEvents map[int][]*agentEventWrap
}
type workflowLoopState struct {
	Iter int
	Idx  int
}

type agentEventWrap struct{ Event any }

type WorkflowInterruptInfo struct {
	OrigInput                *AgentInput
	SequentialIdx            int
	SequentialInfo           *InterruptInfo
	LoopIter                 int
	ParallelInfo             map[int]*InterruptInfo
}

type workflowAgent struct {
	name          string
	desc          string
	subAgents     []*flowAgent
	mode          workflowMode
	maxIter       int
}

func (a *workflowAgent) Name(_ context.Context) string               { return a.name }
func (a *workflowAgent) Description(_ context.Context) string        { return a.desc }
func (a *workflowAgent) GetType() string {
	switch a.mode {
	case workflowModeSequential: return "Sequential"
	case workflowModeParallel:   return "Parallel"
	case workflowModeLoop:       return "Loop"
	default:               return "WorkflowAgent"
	}
}

func (a *workflowAgent) Run(ctx context.Context, _ *AgentInput, opts ...RunOption) *AsyncIterator[*AgentEvent] {
	it, gen := NewAsyncIteratorPair[*AgentEvent]()
	go func() {
		defer func() {
			if r := recover(); r != nil { gen.Send(&AgentEvent{Err: fmt.Errorf("panic: %v\n%s", r, debug.Stack())}) }
			gen.Close()
		}()
		switch a.mode {
		case workflowModeSequential: a.runSeq(ctx, gen, nil, nil, opts...)
		case workflowModeParallel:   a.runPar(ctx, gen, nil, nil, opts...)
		case workflowModeLoop:       a.runLoop(ctx, gen, nil, nil, opts...)
		default:               gen.Send(&AgentEvent{Err: fmt.Errorf("unsupported mode %d", a.mode)})
		}
	}()
	return it
}

func (a *workflowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...RunOption) *AsyncIterator[*AgentEvent] {
	it, gen := NewAsyncIteratorPair[*AgentEvent]()
	go func() {
		defer func() {
			if r := recover(); r != nil { gen.Send(&AgentEvent{Err: fmt.Errorf("panic: %v\n%s", r, debug.Stack())}) }
			gen.Close()
		}()
		st := info.InterruptState
		if st == nil { gen.Send(&AgentEvent{Err: fmt.Errorf("no state for resume")}); return }
		switch s := st.(type) {
		case *workflowState:         a.runSeq(ctx, gen, s, info, opts...)
		case *workflowParallelState: a.runPar(ctx, gen, s, info, opts...)
		case *workflowLoopState:     a.runLoop(ctx, gen, s, info, opts...)
		default:               gen.Send(&AgentEvent{Err: fmt.Errorf("unknown state %T", s)})
		}
	}()
	return it
}

// ---- Sequential ----

func (a *workflowAgent) runSeq(ctx context.Context, gen *AsyncGenerator[*AgentEvent], st *workflowState, info *ResumeInfo, opts ...RunOption) error {
	start := 0
	wfCtx := ctx
	if st != nil { start = st.InterruptIdx; wfCtx = buildPath(ctx, a.subAgents, start, 0) }

	for i := start; i < len(a.subAgents); i++ {
		sa := a.subAgents[i]
		if cc := getCancelContext(ctx); cc != nil && cc.shouldCancel() {
			gen.Send(cancelTransition(ctx, "Sequential cancel", &workflowState{InterruptIdx: i})); return nil
		}
		var si *AsyncIterator[*AgentEvent]
		if st != nil {
			if wfInfo, _ := info.Data.(*WorkflowInterruptInfo); wfInfo != nil && wfInfo.SequentialInfo != nil {
				si = sa.Resume(wfCtx, &ResumeInfo{EnableStreaming: info.EnableStreaming, InterruptInfo: wfInfo.SequentialInfo}, opts...)
			} else { si = sa.Run(wfCtx, nil, opts...) }
			st = nil
		} else { si = sa.Run(wfCtx, nil, opts...) }

		wfCtx = updateRunPathOnly(wfCtx, sa.Name(wfCtx))
		last := drainEvents(si, gen)
		if last != nil {
			if last.Err != nil {
				gen.Send(last); return nil
			}
			if last.Action.internalInterrupted != nil {
				s := &workflowState{InterruptIdx: i}
				ev := CompositeInterrupt(ctx, "Seq interrupted", s, last.Action.internalInterrupted)
				ev.Action.Interrupted.Data = &WorkflowInterruptInfo{OrigInput: inputFromCtx(ctx), SequentialIdx: i, SequentialInfo: last.Action.Interrupted}
				ev.AgentName, ev.RunPath = last.AgentName, last.RunPath
				gen.Send(ev); return nil
			}
			if last.Action.Exit { gen.Send(last); return nil }
			gen.Send(last)
		}
	}
	return nil
}

// ---- Loop ----

func (a *workflowAgent) runLoop(ctx context.Context, gen *AsyncGenerator[*AgentEvent], ls *workflowLoopState, info *ResumeInfo, opts ...RunOption) error {
	if len(a.subAgents) == 0 { return nil }
	startIter, startIdx := 0, 0
	wfCtx := ctx
	if ls != nil { startIter, startIdx = ls.Iter, ls.Idx; wfCtx = buildPath(ctx, a.subAgents, startIdx, startIter) }

	for i := startIter; i < a.maxIter || a.maxIter == 0; i++ {
		for j := startIdx; j < len(a.subAgents); j++ {
			sa := a.subAgents[j]
			if cc := getCancelContext(ctx); cc != nil && cc.shouldCancel() {
				gen.Send(cancelTransition(ctx, "Loop cancel", &workflowLoopState{Iter: i, Idx: j})); return nil
			}
			var si *AsyncIterator[*AgentEvent]
			if ls != nil {
				if wfInfo, _ := info.Data.(*WorkflowInterruptInfo); wfInfo != nil && wfInfo.SequentialInfo != nil {
					si = sa.Resume(wfCtx, &ResumeInfo{EnableStreaming: info.EnableStreaming, InterruptInfo: wfInfo.SequentialInfo}, opts...)
				} else { si = sa.Run(wfCtx, nil, opts...) }
				ls = nil
			} else { si = sa.Run(wfCtx, nil, opts...) }

			wfCtx = updateRunPathOnly(wfCtx, sa.Name(wfCtx))
			var breakEv *AgentEvent
			_ = breakEv
			last := drainEvents(si, gen)
			if last != nil {
				if last.Err != nil {
					gen.Send(last); return nil
				}
				if last.Action.BreakLoop != nil && !last.Action.BreakLoop.Done {
					last.Action.BreakLoop.Done = true
					last.Action.BreakLoop.CurrentIterations = i
					gen.Send(last)
					return nil
				}
				if last.Action.internalInterrupted != nil {
					s := &workflowLoopState{Iter: i, Idx: j}
					ev := CompositeInterrupt(ctx, "Loop interrupted", s, last.Action.internalInterrupted)
					ev.Action.Interrupted.Data = &WorkflowInterruptInfo{OrigInput: inputFromCtx(ctx), LoopIter: i, SequentialIdx: j, SequentialInfo: last.Action.Interrupted}
					ev.AgentName, ev.RunPath = last.AgentName, last.RunPath
					gen.Send(ev); return nil
				}
				if last.Action.Exit { gen.Send(last); return nil }
				gen.Send(last)
			}
		}
		startIdx = 0
	}
	return nil
}

// ---- Parallel ----

func (a *workflowAgent) runPar(ctx context.Context, gen *AsyncGenerator[*AgentEvent], ps *workflowParallelState, info *ResumeInfo, opts ...RunOption) error {
	if len(a.subAgents) == 0 { return nil }
	var wg sync.WaitGroup
	var mu sync.Mutex
	var signals []*InterruptSignal
	dataMap := make(map[int]*InterruptInfo)
	var names map[string]bool

	if ps != nil {
		n, err := getNextResumeAgents(ctx, info)
		if err != nil { return err }
		names = n
	}
	childCtxs := make([]context.Context, len(a.subAgents))
	for i := range a.subAgents {
		childCtxs[i] = forkRunCtx(ctx)
		if ps != nil && ps.SubEvents != nil {
			if evts, ok := ps.SubEvents[i]; ok {
				if rc := getRunCtx(childCtxs[i]); rc != nil && rc.Session != nil {
					for _, e := range evts { rc.Session.addEvent(e) }
				}
			}
		}
	}
	if cc := getCancelContext(ctx); cc != nil && cc.shouldCancel() {
		gen.Send(cancelTransition(ctx, "Parallel cancel", &workflowParallelState{})); return nil
	}
	for i := range a.subAgents {
		wg.Add(1)
		go func(idx int, ag *flowAgent) {
			defer wg.Done()
			var it *AsyncIterator[*AgentEvent]
			if names != nil {
				if _, ok := names[ag.Name(ctx)]; ok {
					ri := &ResumeInfo{EnableStreaming: info.EnableStreaming}
					if wf, _ := info.Data.(*WorkflowInterruptInfo); wf != nil { ri.InterruptInfo = wf.ParallelInfo[idx] }
					it = ag.Resume(childCtxs[idx], ri, opts...)
				} else if ps != nil {
					return
				} else {
					it = ag.Run(childCtxs[idx], nil, opts...)
				}
			} else { it = ag.Run(childCtxs[idx], nil, opts...) }

			for { ev, ok := it.Next(); if !ok { break }
				if ev.Action != nil && ev.Action.internalInterrupted != nil {
					mu.Lock(); signals = append(signals, ev.Action.internalInterrupted); dataMap[idx] = ev.Action.Interrupted; mu.Unlock(); break
				}
				gen.Send(ev)
			}
		}(i, a.subAgents[i])
	}
	wg.Wait()
	if len(signals) > 0 {
		subEvts := make(map[int][]*agentEventWrap)
		for i, cc := range childCtxs {
			if rc := getRunCtx(cc); rc != nil && rc.Session != nil {
				var ws []*agentEventWrap
				for _, e := range rc.Session.getEvents() { ws = append(ws, &agentEventWrap{Event: e}) }
				subEvts[i] = ws
			}
		}
		st := &workflowParallelState{SubEvents: subEvts}
		ev := CompositeInterrupt(ctx, "Parallel interrupted", st, signals...)
		ev.Action.Interrupted.Data = &WorkflowInterruptInfo{OrigInput: inputFromCtx(ctx), ParallelInfo: dataMap}
		ev.AgentName = a.Name(ctx); ev.RunPath = getRunCtx(ctx).getRunPath()
		gen.Send(ev)
	}
	return nil
}

// ---- Helpers ----

func buildPath(ctx context.Context, subs []*flowAgent, idx, iter int) context.Context {
	var steps []string
	for k := 0; k < iter; k++ { for _, s := range subs { steps = append(steps, s.Name(ctx)) } }
	for k := 0; k < idx; k++ { steps = append(steps, subs[k].Name(ctx)) }
	return updateRunPathOnly(ctx, steps...)
}

func drainEvents(ai *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent]) *AgentEvent {
	var last *AgentEvent
	for { ev, ok := ai.Next(); if !ok { break }
		if ev.Err != nil {
			// Return error event instead of sending it to gen — caller handles propagation.
			return ev
		}
		if ev.Action != nil { last = ev; continue }
		gen.Send(ev)
	}
	return last
}

func cancelTransition(ctx context.Context, msg string, state any) *AgentEvent {
	return &AgentEvent{Action: &AgentAction{Interrupted: &InterruptInfo{Data: msg}, internalInterrupted: &InterruptSignal{Info: msg, State: state}}}
}

func inputFromCtx(ctx context.Context) *AgentInput {
	if rc := getRunCtx(ctx); rc != nil { if in, ok := rc.RootInput.(*AgentInput); ok { return in } }
	return nil
}

// ---- Constructors ----

type SequentialConfig struct{ Name, Description string; SubAgents []Agent }
type ParallelConfig struct{ Name, Description string; SubAgents []Agent }
type LoopConfig struct{ Name, Description string; SubAgents []Agent; MaxIterations int }

func newWf(ctx context.Context, name, desc string, subs []Agent, mode workflowMode, maxIter int) (*flowAgent, error) {
	wa := &workflowAgent{name: name, desc: desc, mode: mode, maxIter: maxIter}
	fas := make([]Agent, len(subs))
	for i, s := range subs { fas[i] = toFlowAgent(ctx, s, WithDisallowTransferToParent()) }
	fa, err := SetSubAgents(ctx, wa, fas)
	if err != nil { return nil, err }
	// Set sub-agents directly on the workflowAgent so its Run() has access
	wa.subAgents = make([]*flowAgent, len(fas))
	for i, s := range fas {
		wa.subAgents[i] = toFlowAgent(ctx, s, WithDisallowTransferToParent())
	}
	return fa.(*flowAgent), nil
}

func NewSequential(ctx context.Context, cfg *SequentialConfig) (ResumableAgent, error) {
	if cfg == nil { return nil, fmt.Errorf("SequentialConfig is nil") }
	return newWf(ctx, cfg.Name, cfg.Description, cfg.SubAgents, workflowModeSequential, 0)
}
func NewParallel(ctx context.Context, cfg *ParallelConfig) (ResumableAgent, error) {
	if cfg == nil { return nil, fmt.Errorf("ParallelConfig is nil") }
	return newWf(ctx, cfg.Name, cfg.Description, cfg.SubAgents, workflowModeParallel, 0)
}
func NewLoop(ctx context.Context, cfg *LoopConfig) (ResumableAgent, error) {
	if cfg == nil { return nil, fmt.Errorf("LoopConfig is nil") }
	if cfg.MaxIterations <= 0 { cfg.MaxIterations = 10 }
	return newWf(ctx, cfg.Name, cfg.Description, cfg.SubAgents, workflowModeLoop, cfg.MaxIterations)
}

func init() {
	schema.RegisterType("_harness_wf_interrupt_info", func() any { return &WorkflowInterruptInfo{} })
	schema.RegisterType("_harness_wf_state", func() any { return &workflowState{} })
	schema.RegisterType("_harness_wf_parallel_state", func() any { return &workflowParallelState{} })
	schema.RegisterType("_harness_wf_loop_state", func() any { return &workflowLoopState{} })
}