Agent-Native/runtime/session_agent_test.go

package runtime

import (
	"context"
	"sync"
	"testing"

	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/llm"
)

func TestAgentSessionPromptForwardsStop(t *testing.T) {
	called := false
	sess := NewAgentSession(Config{
		Model: "m",
		Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "hi", StopReason: llm.StopEnd}}
			close(ch)
			return ch
		},
		Stop: func(last llm.Message, info StopInfo) StopDecision {
			called = true
			return StopDecision{}
		},
	})
	res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if !called || res.Text != "hi" {
		t.Fatalf("called=%v text=%q", called, res.Text)
	}
}

func TestAgentSessionSteerFollowAndAbort(t *testing.T) {
	var mu sync.Mutex
	step := 0
	stream := llm.StreamFunc(func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 2)
		go func() {
			defer close(ch)
			mu.Lock()
			n := step
			step++
			mu.Unlock()
			if opts.ThinkingLevel != "low" && n == 0 {
				t.Errorf("thinking level not forwarded: %s", opts.ThinkingLevel)
			}
			if n == 0 {
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
					Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
					ToolCalls: []llm.ToolCall{{ID: "1", Name: "missing", Arguments: "{}"}},
				}}
				return
			}
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
				Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd,
			}}
		}()
		return ch
	})
	var types []string
	sess := NewAgentSession(Config{
		Model:         "m",
		ThinkingLevel: "low",
		Stream:        stream,
		MaxTurns:      6,
		OnEvent: func(ev Event) {
			types = append(types, ev.Type)
		},
	})
	sess.SetQueueMode(DeliverOne, DeliverAll)
	sess.Steer(llm.Message{Role: llm.RoleUser, Content: "steer-me"})
	sess.FollowUp(llm.Message{Role: llm.RoleUser, Content: "later"})
	res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.StopReason != llm.StopEnd {
		t.Fatalf("stop %s err %v", res.StopReason, res.Err)
	}
	found := map[string]bool{}
	for _, typ := range types {
		found[typ] = true
	}
	for _, want := range []string{"message_start", "queue_update", "tool_execution_update", "agent_settled"} {
		if !found[want] {
			t.Fatalf("missing %s in %v", want, types)
		}
	}
}

func TestClearQueueAndSetModel(t *testing.T) {
	sess := NewAgentSession(Config{Model: "old", Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 1)
		ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: model, StopReason: llm.StopEnd}}
		close(ch)
		return ch
	}})
	sess.Steer(llm.Message{Role: llm.RoleUser, Content: "drop"})
	sess.ClearQueue()
	if err := sess.SetModel("new-model", "high"); err != nil {
		t.Fatal(err)
	}
	res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "hi"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Text != "new-model" {
		t.Fatalf("model switch did not apply: %q", res.Text)
	}
}