Agent-Native/runtime/llm/fallback_test.go

package llm

import (
	"context"
	"testing"
)

func TestFallbackSwitchesOn5xxBeforeOutput(t *testing.T) {
	var switched string
	base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
		ch := make(chan Event, 2)
		go func() {
			defer close(ch)
			if model == "primary" {
				ch <- Event{Kind: KindError, HTTPStatus: 503, ErrorMessage: "provider status 503", Message: Message{Role: RoleAssistant, StopReason: StopError}}
				return
			}
			ch <- Event{Kind: KindTextDelta, Text: "ok"}
			ch <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok", StopReason: StopEnd}}
		}()
		return ch
	}
	stream := WithFallback(base, []string{"backup"}, func(from, to string) {
		switched = from + "->" + to
	})
	var text string
	var sawError bool
	for ev := range stream(context.Background(), "primary", nil, nil, StreamOptions{}) {
		if ev.Kind == KindTextDelta {
			text += ev.Text
		}
		if ev.Kind == KindError {
			sawError = true
		}
	}
	if switched != "primary->backup" || text != "ok" || sawError {
		t.Fatalf("switched %q text %q error %v", switched, text, sawError)
	}
}

func TestFallbackStaysAfterFirstDelta(t *testing.T) {
	calls := 0
	base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
		calls++
		ch := make(chan Event, 2)
		go func() {
			defer close(ch)
			ch <- Event{Kind: KindTextDelta, Text: "partial"}
			ch <- Event{Kind: KindError, HTTPStatus: 500, Message: Message{Role: RoleAssistant, StopReason: StopError}}
		}()
		return ch
	}
	stream := WithFallback(base, []string{"backup"}, func(from, to string) {
		t.Fatal("switched after output")
	})
	n := 0
	for range stream(context.Background(), "primary", nil, nil, StreamOptions{}) {
		n++
	}
	if calls != 1 || n != 2 {
		t.Fatalf("calls %d events %d", calls, n)
	}
}

func TestEmptyFallbackIsIdentity(t *testing.T) {
	base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
		ch := make(chan Event)
		close(ch)
		return ch
	}
	if WithFallback(base, nil, nil) == nil {
		t.Fatal("nil stream")
	}
}