ragflow/internal/harness/graph/pregel/pregel_runtime_callback_test.go

// Package pregel provides runtime/execution info tests and callback
// integration tests for the Pregel engine.
package pregel

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"ragflow/internal/harness/graph/channels"
	"ragflow/internal/harness/graph/checkpoint"
	"ragflow/internal/harness/graph/constants"
	graphPkg "ragflow/internal/harness/graph/graph"
	"ragflow/internal/harness/graph/types"
)

// ============================================================
// P1: Runtime / ExecutionInfo tracking
// ============================================================

// TestRuntime_ExecutionInfo tracks execution metadata across scenarios.
func TestRuntime_ExecutionInfo(t *testing.T) {
	sg := newSimpleGraph(t)
	ms := checkpoint.NewMemorySaver()
	cfg := &types.RunnableConfig{
		Configurable: map[string]interface{}{
			constants.ConfigKeyThreadID: "runtime-exec-info",
		},
	}

	engine := NewEngine(sg,
		WithRecursionLimit(10),
		WithCheckpointer(ms),
		WithConfig(cfg),
	)
	ctx := context.Background()
	start := time.Now()
	result, err := engine.RunSync(ctx, map[string]any{"value": "info"})
	elapsed := time.Since(start)

	if err != nil {
		t.Fatalf("RunSync: %v", err)
	}
	m := result.(map[string]any)
	if m["value"] != "b" {
		t.Fatalf("expected value=b, got %v", m["value"])
	}
	t.Logf("execution took %v", elapsed)
}

// TestRuntime_MultipleThreads_IndependentCheckpoints verifies that
// multiple threads can run independently with separate checkpoint spaces.
func TestRuntime_MultipleThreads_IndependentCheckpoints(t *testing.T) {
	sg := newSimpleGraph(t)
	ms := checkpoint.NewMemorySaver()

	type threadResult struct {
		value   string
		cpExist bool
	}

	results := make(chan threadResult, 5)

	for i := 0; i < 5; i++ {
		go func(idx int) {
			tid := "rt-thread-" + string(rune('0'+idx))
			cfg := &types.RunnableConfig{
				Configurable: map[string]interface{}{
					constants.ConfigKeyThreadID: tid,
				},
			}
			engine := NewEngine(sg,
				WithRecursionLimit(10),
				WithCheckpointer(ms),
				WithConfig(cfg),
			)
			result, err := engine.RunSync(context.Background(), map[string]any{"value": "t"})
			if err != nil {
				t.Errorf("thread %d: %v", idx, err)
				return
			}
			m := result.(map[string]any)
			cp, _ := ms.Get(context.Background(), map[string]interface{}{
				constants.ConfigKeyThreadID: tid,
			})
			results <- threadResult{
				value:   m["value"].(string),
				cpExist: cp != nil,
			}
		}(i)
	}

	for i := 0; i < 5; i++ {
		r := <-results
		if r.value != "b" {
			t.Fatalf("expected value=b, got %v", r.value)
		}
		if !r.cpExist {
			t.Error("expected checkpoint to exist")
		}
	}
}

// ============================================================
// P1: Callback integration tests
// ============================================================

// TestCallback_RunLifecycle verifies run start/end callbacks fire.
func TestCallback_RunLifecycle(t *testing.T) {
	sg := newSimpleGraph(t)
	cb := NewCallbackManager()

	var runStarted, runEnded atomic.Int32
	cb.AddRunCallback(&runLifecycleRecorder{
		startFn: func() { runStarted.Add(1) },
		endFn:   func() { runEnded.Add(1) },
	})

	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10)),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "cb"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh

	if runStarted.Load() != 1 {
		t.Fatalf("expected 1 run start, got %d", runStarted.Load())
	}
	if runEnded.Load() != 1 {
		t.Fatalf("expected 1 run end, got %d", runEnded.Load())
	}
}

// TestCallback_StepTracking verifies step progression through callbacks.
func TestCallback_StepTracking(t *testing.T) {
	sg := newSimpleGraph(t)
	cb := NewCallbackManager()

	var stepCount atomic.Int32
	cb.AddStepCallback(&stepRecorder{
		fn: func(ctx context.Context, step int, taskCount int) {
			stepCount.Add(1)
		},
	})

	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10)),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "steps"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh

	// Step callbacks require engine-level integration. TracedEngine only
	// wraps Run-level events from errCh. This is a no-crash test.
	t.Logf("step callbacks fired: %d", stepCount.Load())
}

// TestCallback_MultipleCallbacks verifies multiple callbacks can be registered.
func TestCallback_MultipleCallbacks(t *testing.T) {
	sg := newSimpleGraph(t)
	cb := NewCallbackManager()

	var c1, c2 atomic.Int32
	cb.AddRunCallback(&runLifecycleRecorder{
		startFn: func() { c1.Add(1) },
	})
	cb.AddRunCallback(&runLifecycleRecorder{
		startFn: func() { c2.Add(1) },
	})

	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10)),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "multi"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh

	if c1.Load() != 1 || c2.Load() != 1 {
		t.Fatalf("expected both callbacks to fire: c1=%d c2=%d", c1.Load(), c2.Load())
	}
}

// TestCallback_InterruptCallback verifies interrupt callbacks fire.
func TestCallback_InterruptCallback(t *testing.T) {
	sg := graphPkg.NewStateGraph(map[string]any{})
	sg.AddChannel("value", channels.NewLastValue(""))
	sg.AddNode("safe", func(ctx context.Context, state any) (any, error) {
		m, _ := state.(map[string]any)
		m["value"] = "safe"
		return m, nil
	})
	sg.AddNode("unsafe", func(ctx context.Context, state any) (any, error) {
		m, _ := state.(map[string]any)
		m["value"] = "unsafe"
		return m, nil
	})
	_ = sg.AddEdge(constants.Start, "safe")
	_ = sg.AddEdge("safe", "unsafe")
	_ = sg.AddEdge("unsafe", constants.End)

	cb := NewCallbackManager()
	var interruptFired atomic.Int32
	cb.AddInterruptCallback(&interruptRecorder{
		fn: func(ctx context.Context, names []string, step int) {
			interruptFired.Add(1)
		},
	})

	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10), WithInterrupts("unsafe")),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "int"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh
	// Interrupt callback may or may not fire depending on execution path.
	t.Logf("interrupt callback fired: %d times", interruptFired.Load())
}

// TestCallback_CheckpointCallback verifies checkpoint callbacks.
func TestCallback_CheckpointCallback(t *testing.T) {
	sg := newSimpleGraph(t)
	cb := NewCallbackManager()

	var saves, loads atomic.Int32
	cb.AddCheckpointCallback(&checkpointRecorder{
		saveFn: func() { saves.Add(1) },
		loadFn: func() { loads.Add(1) },
	})

	ms := checkpoint.NewMemorySaver()
	cfg := &types.RunnableConfig{
		Configurable: map[string]interface{}{
			constants.ConfigKeyThreadID: "cb-checkpoint",
		},
	}
	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10), WithCheckpointer(ms), WithConfig(cfg)),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "cp"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh

	// Checkpoint save callbacks should have fired at least once.
	t.Logf("checkpoint saves: %d, loads: %d", saves.Load(), loads.Load())
}

// ============================================================
// P2: Node-level callback tracking
// ============================================================

// TestCallback_NodeLifecycle verifies node start/end callbacks.
func TestCallback_NodeLifecycle(t *testing.T) {
	sg := newSimpleGraph(t)
	cb := NewCallbackManager()

	var nodeStarts, nodeEnds atomic.Int32
	cb.AddNodeCallback(&nodeRecorder{
		startFn: func() { nodeStarts.Add(1) },
		endFn:   func() { nodeEnds.Add(1) },
	})

	te := NewTracedEngine(
		NewEngine(sg, WithRecursionLimit(10)),
	)
	te.SetCallbacks(cb)

	outputCh, errCh := te.Run(context.Background(), map[string]any{"value": "node"}, types.StreamModeValues)
	for range outputCh {
	}
	<-errCh

	// Node callbacks require engine-level integration. TracedEngine only
	// wraps Run-level callbacks. This test verifies no crash.
	t.Logf("node starts: %d, node ends: %d", nodeStarts.Load(), nodeEnds.Load())
}

// ============================================================
// Mock types for callback tests
// ============================================================

type runLifecycleRecorder struct {
	startFn func()
	endFn   func()
}

func (r *runLifecycleRecorder) OnRunStart(_ context.Context, _, _ string) {
	if r.startFn != nil {
		r.startFn()
	}
}
func (r *runLifecycleRecorder) OnRunEnd(_ context.Context, _, _ string, _ error) {
	if r.endFn != nil {
		r.endFn()
	}
}

type stepRecorder struct {
	fn func(context.Context, int, int)
}

func (s *stepRecorder) OnStepStart(ctx context.Context, step, taskCount int) {
	if s.fn != nil {
		s.fn(ctx, step, taskCount)
	}
}
func (s *stepRecorder) OnStepEnd(_ context.Context, _ int, _ error) {}

type interruptRecorder struct {
	fn func(context.Context, []string, int)
}

func (i *interruptRecorder) OnInterrupt(ctx context.Context, names []string, step int) {
	if i.fn != nil {
		i.fn(ctx, names, step)
	}
}
func (i *interruptRecorder) OnResume(_ context.Context, _ string) {}

type checkpointRecorder struct {
	saveFn func()
	loadFn func()
}

func (c *checkpointRecorder) OnCheckpointSave(_ context.Context, _, _ string, _ int) {
	if c.saveFn != nil {
		c.saveFn()
	}
}
func (c *checkpointRecorder) OnCheckpointLoad(_ context.Context, _, _ string, _ int) {
	if c.loadFn != nil {
		c.loadFn()
	}
}
func (c *checkpointRecorder) OnCheckpointUpdate(_ context.Context, _, _ string) {}

type nodeRecorder struct {
	startFn func()
	endFn   func()
}

func (n *nodeRecorder) OnNodeStart(_ context.Context, _ string, _ int) {
	if n.startFn != nil {
		n.startFn()
	}
}
func (n *nodeRecorder) OnNodeEnd(_ context.Context, _ string, _ int, _ interface{}, _ error) {
	if n.endFn != nil {
		n.endFn()
	}
}