Agent-Native/runtime/loop_harness_test.go

package runtime

import (
	"context"
	"encoding/json"
	"errors"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"agent-native/runtime/llm"
	"agent-native/runtime/session"
	"agent-native/runtime/tools"
)

func TestLoopStopContinuesUntilCap(t *testing.T) {
	var calls int
	res, err := Run(context.Background(), Config{
		Model:    "fake",
		MaxTurns: 8,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
		),
		Stop: func(last llm.Message, info StopInfo) StopDecision {
			calls++
			if info.Continuations != calls-1 {
				t.Errorf("continuations = %d, want %d", info.Continuations, calls-1)
			}
			return StopDecision{Continue: true, Replacement: "已作废", Reason: "再试", Notice: "到顶了"}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if calls != 4 {
		t.Fatalf("stop calls = %d, want 4", calls)
	}
	if res.Err != nil || res.Notice != "到顶了" {
		t.Fatalf("notice %q err %v", res.Notice, res.Err)
	}
	replaced, forged := 0, 0
	for _, m := range res.Messages {
		if m.Role != llm.RoleAssistant {
			continue
		}
		if m.Content == "已作废" {
			replaced++
		}
		if m.Content == "forged-claim" {
			forged++
		}
	}
	if replaced != 3 || forged != 1 {
		t.Fatalf("replaced=%d forged=%d transcript=%s", replaced, forged, transcriptText(res.Messages))
	}
	if strings.Count(transcriptText(res.Messages), "再试") != 3 {
		t.Fatalf("reasons = %s", transcriptText(res.Messages))
	}
}

func TestLoopBeforeToolRejectPairsResult(t *testing.T) {
	hits := 0
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(hitTool{hits: &hits})
	var seen llm.Message
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Tools: reg,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "echo", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
		BeforeTool: func(call llm.ToolCall) error {
			return errors.New("不允许")
		},
		AfterTool: func(call llm.ToolCall, msg llm.Message) {
			seen = msg
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if hits != 0 {
		t.Fatalf("tool executed %d times", hits)
	}
	if !seen.IsError || seen.Content != "不允许" || seen.ToolCallID != "c1" {
		t.Fatalf("after = %#v", seen)
	}
	paired := false
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError && m.Content == "不允许" {
			paired = true
		}
	}
	if !paired {
		t.Fatalf("missing paired tool result: %s", transcriptText(res.Messages))
	}
}

func TestLoopRemindersStayOutOfSession(t *testing.T) {
	path := filepath.Join(t.TempDir(), "transcript.jsonl")
	st, err := session.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	var saw bool
	_, err = Run(context.Background(), Config{
		Model:   "fake",
		Session: st,
		Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			for _, m := range messages {
				if strings.Contains(m.Content, "<system-reminder>") && strings.Contains(m.Content, "先核对交付物") {
					saw = true
				}
			}
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
			close(ch)
			return ch
		},
		Reminders: func(turn int) []string {
			if turn != 0 {
				t.Errorf("turn %d", turn)
			}
			return []string{"先核对交付物"}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "ping"})
	if err != nil {
		t.Fatal(err)
	}
	if !saw {
		t.Fatal("reminder missing from the provider request")
	}
	loaded, err := st.Load()
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(transcriptText(loaded), "先核对交付物") {
		t.Fatalf("reminder persisted: %#v", loaded)
	}
}

func TestLoopRepeatGuardFiresAndCaps(t *testing.T) {
	var mu sync.Mutex
	var requests []string
	turns := make([]llm.Message, 0, 13)
	for i := 0; i < 12; i++ {
		args := `{"b":1,"a":2}`
		if i%2 == 1 {
			args = `{"a":2,"b":1}`
		}
		turns = append(turns, llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "c", Name: "echo", Arguments: args,
		}}})
	}
	turns = append(turns, llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd})
	var i int
	stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		mu.Lock()
		requests = append(requests, transcriptText(messages))
		n := i
		i++
		mu.Unlock()
		ch := make(chan llm.Event, 1)
		msg := turns[len(turns)-1]
		if n < len(turns) {
			msg = turns[n]
		}
		ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
		close(ch)
		return ch
	}
	res, err := Run(context.Background(), Config{
		Model:       "fake",
		MaxTurns:    20,
		RepeatGuard: true,
		Stream:      stream,
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Text != "done" {
		t.Fatalf("text %q", res.Text)
	}
	reminders := 0
	for _, req := range requests {
		reminders += strings.Count(req, repeatReminderText)
	}
	if reminders != 3 {
		t.Fatalf("reminders = %d, want 3; requests=%d", reminders, len(requests))
	}
	if strings.Contains(transcriptText(res.Messages), repeatReminderText) {
		t.Fatal("repeat reminder was written into the transcript")
	}
	if !strings.Contains(requests[3], repeatReminderText) {
		t.Fatal("first reminder should arrive on the request after the third identical call")
	}
}

func TestLoopContinueOnLengthCaps(t *testing.T) {
	var requests []string
	var n int
	stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		requests = append(requests, transcriptText(messages))
		n++
		ch := make(chan llm.Event, 1)
		ch <- llm.Event{
			Kind:         llm.KindMessageEnd,
			FinishReason: "length",
			Message:      llm.Message{Role: llm.RoleAssistant, Content: "part", StopReason: llm.StopEnd},
		}
		close(ch)
		return ch
	}
	res, err := Run(context.Background(), Config{
		Model:            "fake",
		MaxTurns:         8,
		ContinueOnLength: true,
		Stream:           stream,
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || !strings.Contains(res.Err.Error(), "length continuation exceeded") {
		t.Fatalf("err = %v", res.Err)
	}
	if n != 4 {
		t.Fatalf("model calls = %d, want 4", n)
	}
	if len(res.Messages) == 0 || !strings.Contains(transcriptText(res.Messages), "part") {
		t.Fatal("truncated transcript was dropped")
	}
	if strings.Contains(transcriptText(res.Messages), lengthContinuePrompt) {
		t.Fatal("continuation prompt was persisted")
	}
	prompts := 0
	for _, req := range requests {
		prompts += strings.Count(req, lengthContinuePrompt)
	}
	if prompts != 3 {
		t.Fatalf("continuation prompts = %d, want 3", prompts)
	}
}

func TestLoopContinueOnLengthUnlimited(t *testing.T) {
	var n int
	stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		n++
		ch := make(chan llm.Event, 1)
		finish := "length"
		text := "part"
		if n >= 5 {
			finish = "stop"
			text = "done"
		}
		ch <- llm.Event{
			Kind:         llm.KindMessageEnd,
			FinishReason: finish,
			Message:      llm.Message{Role: llm.RoleAssistant, Content: text, StopReason: llm.StopEnd},
		}
		close(ch)
		return ch
	}
	res, err := Run(context.Background(), Config{
		Model:                  "fake",
		MaxTurns:               8,
		ContinueOnLength:       true,
		MaxLengthContinuations: -1,
		Stream:                 stream,
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err != nil {
		t.Fatalf("err = %v", res.Err)
	}
	if n != 5 {
		t.Fatalf("model calls = %d, want 5", n)
	}
	if res.Text != "done" {
		t.Fatalf("text = %q", res.Text)
	}
}

func TestLoopHistorySkipsGeneratedSystem(t *testing.T) {
	var first string
	_, err := Run(context.Background(), Config{
		Model:        "fake",
		SystemPrompt: "LEAK",
		History: []llm.Message{
			{Role: llm.RoleSystem, Content: "FROM_HISTORY"},
			{Role: llm.RoleUser, Content: "earlier"},
		},
		Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			if len(messages) > 0 {
				first = messages[0].Content
			}
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
			close(ch)
			return ch
		},
	}, llm.Message{Role: llm.RoleUser, Content: "now"})
	if err != nil {
		t.Fatal(err)
	}
	if first != "FROM_HISTORY" || strings.Contains(first, "LEAK") {
		t.Fatalf("system = %q", first)
	}
}

func TestLoopUnknownToolFallsThrough(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model: "fake",
		UnknownTool: func(call llm.ToolCall) tools.Result {
			return tools.Result{Content: "via " + call.Name}
		},
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "make_report", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	found := false
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.Content == "via make_report" && !m.IsError {
			found = true
		}
	}
	if !found {
		t.Fatalf("unknown tool was not delegated: %s", transcriptText(res.Messages))
	}
}

func TestLoopTurnCapKeepsTranscript(t *testing.T) {
	res, err := Run(context.Background(), Config{
		Model:    "fake",
		MaxTurns: 1,
		Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "c1", Name: "missing", Arguments: `{}`,
		}}}),
	}, llm.Message{Role: llm.RoleUser, Content: "ping"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || !strings.Contains(res.Err.Error(), "max turns exceeded") {
		t.Fatalf("err = %v", res.Err)
	}
	if !strings.Contains(transcriptText(res.Messages), "ping") {
		t.Fatalf("transcript dropped: %s", transcriptText(res.Messages))
	}
}

type hitTool struct{ hits *int }

func (h hitTool) Name() string                { return "echo" }
func (h hitTool) Description() string         { return "echo" }
func (h hitTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
func (h hitTool) Execute(context.Context, json.RawMessage) (tools.Result, error) {
	*h.hits++
	return tools.Result{Content: "ran"}, nil
}