ragflow/internal/harness/core/model_chain_test.go

package core

import (
	"context"
	"io"
	"testing"

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

// ---- BuildModelWrapperChain tests ----

func TestBuildModelWrapperChain_NoConfig(t *testing.T) {
	base := &mockModel{}
	base.addResp("raw")
	model := BuildModelWrapperChain(base, nil, DefaultReActConfig[*schema.Message]())
	if model == nil { t.Fatal("nil model") }
	resp, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "raw" { t.Errorf("content = %s", resp.Content) }
}

func TestBuildModelWrapperChain_WithRetry(t *testing.T) {
	base := &mockModel{}
	base.addResp("retry-ok")
	cfg := DefaultReActConfig[*schema.Message]()
	cfg.RetryConfig = &ModelRetryConfig{MaxRetries: 2}
	model := BuildModelWrapperChain(base, nil, cfg)
	if model == nil { t.Fatal("nil model") }
	resp, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "retry-ok" { t.Errorf("content = %s", resp.Content) }
}

func TestBuildModelWrapperChain_WithFailover(t *testing.T) {
	primary := &mockModel{}
	primary.addResp("primary-ok")
	fallback := &mockModel{}
	fallback.addResp("fallback")

	cfg := DefaultReActConfig[*schema.Message]()
	cfg.FailoverConfig = &FailoverConfigMsg{Models: []Model[*schema.Message]{fallback}}
	model := BuildModelWrapperChain(primary, nil, cfg)
	resp, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "primary-ok" { t.Errorf("content = %s", resp.Content) }
}

func TestBuildModelWrapperChain_WithMiddleware(t *testing.T) {
	var wrapCalled bool
	mw := &testMiddleware{}
	mw.wrapModel = func(ctx context.Context, m Model[*schema.Message], mc *ModelContext) (Model[*schema.Message], error) {
		wrapCalled = true
		return m, nil
	}
	base := &mockModel{}
	base.addResp("mw-ok")
	cfg := DefaultReActConfig[*schema.Message]()
	cfg.Middlewares = []ReActMiddleware{mw}
	model := BuildModelWrapperChain(base, nil, cfg)
	resp, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if !wrapCalled { t.Error("middleware WrapModel not called") }
	_ = resp
}

func TestBuildModelWrapperChain_WithFullChain(t *testing.T) {
	var wrapCalled bool
	mw := &testMiddleware{}
	mw.wrapModel = func(ctx context.Context, m Model[*schema.Message], mc *ModelContext) (Model[*schema.Message], error) {
		wrapCalled = true
		return m, nil
	}

	primary := &mockModel{}
	primary.addResp("chain-ok")
	fallback := &mockModel{}
	fallback.addResp("fallback")

	cfg := DefaultReActConfig[*schema.Message]()
	cfg.RetryConfig = &ModelRetryConfig{MaxRetries: 2}
	cfg.FailoverConfig = &FailoverConfigMsg{Models: []Model[*schema.Message]{fallback}}
	cfg.Middlewares = []ReActMiddleware{mw}

	model := BuildModelWrapperChain(primary, nil, cfg)
	resp, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if !wrapCalled { t.Error("middleware WrapModel not called in chain") }
	_ = resp
}

func TestBuildModelWrapperChain_NilMiddleware(t *testing.T) {
	base := &mockModel{}
	base.addResp("nil-mw")
	cfg := DefaultReActConfig[*schema.Message]()
	cfg.Middlewares = []ReActMiddleware{nil, nil}
	model := BuildModelWrapperChain(base, nil, cfg)
	_, err := model.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
}

// ---- eventSenderModelWrapper tests ----

func TestEventSenderModelWrapper_GenerateSendsEvent(t *testing.T) {
	base := &mockModel{}
	base.addResp("event-test")
	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.Message]]()
	ec := &reActExecCtx{generator: gen}
	wrapped := wrapModelWithEventSender(base, ec)

	go func() {
		defer gen.Close()
		resp, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
		if err != nil { t.Errorf("Generate: %v", err) }
		if resp.Content != "event-test" { t.Errorf("content = %s", resp.Content) }
	}()

	ev, ok := it.Next()
	if !ok { t.Fatal("expected event from wrapper") }
	if ev.Output == nil || ev.Output.MessageOutput == nil {
		t.Error("expected message output event")
	}
}

func TestEventSenderModelWrapper_StreamSendsEvent(t *testing.T) {
	base := &mockModel{}
	base.addResp("stream-event")
	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.Message]]()
	ec := &reActExecCtx{generator: gen}
	wrapped := wrapModelWithEventSender(base, ec)

	go func() {
		defer gen.Close()
		s, err := wrapped.Stream(context.Background(), []Message{schema.UserMessage("hi")})
		if err != nil { t.Errorf("Stream: %v", err) }
		for {
			_, err := s.Recv()
			if err == io.EOF { break }
			if err != nil { t.Errorf("Recv: %v", err); return }
		}
	}()

	ev, ok := it.Next()
	if !ok { t.Fatal("expected event from stream wrapper") }
	if ev.Output == nil || ev.Output.MessageOutput == nil {
		t.Error("expected message output event from stream")
	}
}

func TestEventSenderModelWrapper_SuppressEventSend(t *testing.T) {
	base := &mockModel{}
	base.addResp("suppressed")
	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.Message]]()
	ec := &reActExecCtx{generator: gen, suppressEventSend: true}
	wrapped := wrapModelWithEventSender(base, ec)

	done := make(chan struct{})
	go func() {
		defer close(done)
		defer gen.Close()
		_, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
		if err != nil { t.Errorf("Generate: %v", err) }
	}()

	// After Generate returns, check that the channel has no events
	<-done
	// The generator is closed; try reading one item. If suppressEventSend works,
	// the closed empty channel returns a zero value immediately.
	_, ok := it.Next()
	// When the channel is closed and empty, Next() returns (zero, false)
	if ok {
		t.Error("event should be suppressed")
	}
}

func TestEventSenderModelWrapper_NilExecCtx(t *testing.T) {
	base := &mockModel{}
	base.addResp("nil-ec")
	wrapped := wrapModelWithEventSender(base, nil)
	resp, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "nil-ec" { t.Errorf("content = %s", resp.Content) }
}

func TestEventSenderModelWrapper_NilGenerator(t *testing.T) {
	base := &mockModel{}
	base.addResp("nil-gen")
	ec := &reActExecCtx{generator: nil}
	wrapped := wrapModelWithEventSender(base, ec)
	resp, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "nil-gen" { t.Errorf("content = %s", resp.Content) }
}

func TestEventSenderModelWrapper_BindTools(t *testing.T) {
	base := &mockModel{}
	wrapped := wrapModelWithEventSender(base, nil)
	err := wrapped.BindTools([]*schema.ToolInfo{{Name: "test"}})
	if err != nil { t.Fatalf("BindTools: %v", err) }
}

func TestEventSenderModelWrapper_IsNilMessage(t *testing.T) {
	base := &mockModel{}
	base.addResp("")
	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.Message]]()
	ec := &reActExecCtx{generator: gen}
	wrapped := wrapModelWithEventSender(base, ec)

	done := make(chan struct{})
	go func() {
		defer gen.Close()
		_, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
		if err != nil { t.Errorf("Generate: %v", err) }
		close(done)
	}()
	// Empty content might still send the event because it's not nil
	var hasEvent bool
	select {
	case <-done:
	case <-it.ch:
		hasEvent = true
	}
	_ = hasEvent
}



// ---- callbackModelWrapper tests ----

func TestCallbackModelWrapper_Basic(t *testing.T) {
	inner := &mockModel{}
	inner.addResp("cb-ok")
	wrapped := &callbackModelWrapper[*schema.Message]{inner: inner}
	resp, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Generate: %v", err) }
	if resp.Content != "cb-ok" { t.Errorf("content = %s", resp.Content) }
}

func TestCallbackModelWrapper_Stream(t *testing.T) {
	inner := &mockModel{}
	inner.addResp("cb-stream")
	wrapped := &callbackModelWrapper[*schema.Message]{inner: inner}
	s, err := wrapped.Stream(context.Background(), []Message{schema.UserMessage("hi")})
	if err != nil { t.Fatalf("Stream: %v", err) }
	chunk, err := s.Recv()
	if err != nil { t.Fatalf("Recv: %v", err) }
	if chunk.Content != "cb-stream" { t.Errorf("content = %s", chunk.Content) }
}

func TestCallbackModelWrapper_BindTools(t *testing.T) {
	inner := &mockModel{}
	wrapped := &callbackModelWrapper[*schema.Message]{inner: inner}
	err := wrapped.BindTools([]*schema.ToolInfo{{Name: "test"}})
	if err != nil { t.Fatalf("BindTools: %v", err) }
}

// ---- HasUserEventSenderModelWrapper tests ----

func TestHasUserEventSenderModelWrapper_NilSlice(t *testing.T) {
	if HasUserEventSenderModelWrapper[*schema.Message](nil) {
		t.Error("nil should be false")
	}
}

func TestHasUserEventSenderModelWrapper_WithWrapper(t *testing.T) {
	w := NewEventSenderModelWrapper[*schema.Message]()
	handlers := []TypedReActMiddleware[*schema.Message]{w}
	if !HasUserEventSenderModelWrapper(handlers) {
		t.Error("should detect wrapper")
	}
}

func TestHasUserEventSenderModelWrapper_WithoutWrapper(t *testing.T) {
	mw := &testMiddleware{}
	handlers := []TypedReActMiddleware[*schema.Message]{mw}
	if HasUserEventSenderModelWrapper(handlers) {
		t.Error("should not detect non-wrapper")
	}
}


// ---- Wrapper behavior with ExecCtx integration ----

func TestEventSenderModelWrapper_WithExecCtxGenerator(t *testing.T) {
	base := &mockModel{}
	base.addResp("gen-event")
	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.Message]]()
	ec := &reActExecCtx{generator: gen}
	wrapped := wrapModelWithEventSender(base, ec)

	go func() {
		defer gen.Close()
		_, err := wrapped.Generate(context.Background(), []Message{schema.UserMessage("hi")})
		if err != nil { t.Errorf("Generate: %v", err) }
	}()

	// Read event from iterator to verify generator integration works
	ev, ok := it.Next()
	if !ok { t.Fatal("expected event via generator") }
	if ev.Output == nil || ev.Output.MessageOutput == nil {
		t.Error("expected message output event")
	}
}