Agent-Native/runtime/loop_test.go

package runtime

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

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

func TestRunTerminalToolFailureKeepsUsage(t *testing.T) {
	hits, calls := 0, 0
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(hitTool{hits: &hits})
	stream := func(context.Context, string, []llm.Message, []llm.ToolDecl, llm.StreamOptions) <-chan llm.Event {
		calls++
		out := make(chan llm.Event, 1)
		out <- llm.Event{
			Kind:    llm.KindMessageEnd,
			Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError},
			Usage:   llm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17},
			Err:     errors.New("incomplete tool call"),
		}
		close(out)
		return out
	}
	res, err := Run(context.Background(), Config{Model: "fixture", Tools: reg, Stream: stream},
		llm.Message{Role: llm.RoleUser, Content: "organize materials"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || res.StopReason != llm.StopError || res.Usage.TotalTokens != 17 ||
		hits != 0 || calls != 1 {
		t.Fatalf("terminal failure retried, executed a tool, or lost usage: result=%#v hits=%d calls=%d", res, hits, calls)
	}
}

func TestRunTerminalErrorCannotMasqueradeAsSuccess(t *testing.T) {
	stream := func(context.Context, string, []llm.Message, []llm.ToolDecl, llm.StreamOptions) <-chan llm.Event {
		out := make(chan llm.Event, 1)
		out <- llm.Event{Kind: llm.KindMessageEnd,
			Message: llm.Message{Role: llm.RoleAssistant, Content: "incomplete", StopReason: llm.StopEnd},
			Usage:   llm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17},
			Err:     errors.New("terminal failed"),
		}
		close(out)
		return out
	}
	res, err := Run(context.Background(), Config{Model: "fixture", Stream: stream},
		llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Err == nil || res.StopReason != llm.StopError || res.Usage.TotalTokens != 17 {
		t.Fatalf("error-bearing terminal was accepted as success: %#v", res)
	}
}

func script(turns ...llm.Message) llm.StreamFunc {
	var mu sync.Mutex
	i := 0
	return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 4)
		go func() {
			defer close(ch)
			mu.Lock()
			defer mu.Unlock()
			if i >= len(turns) {
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}}
				return
			}
			msg := turns[i]
			i++
			if msg.Content != "" {
				ch <- llm.Event{Kind: llm.KindTextDelta, Text: msg.Content}
			}
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
		}()
		return ch
	}
}

func TestLoopToolThenContinueAfterToolError(t *testing.T) {
	dir := t.TempDir()
	reg := tools.WorkspaceTools(dir)
	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: "read_file", Arguments: `{"path":"missing.md"}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "recovered", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "read it"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Text != "recovered" || res.StopReason != llm.StopEnd {
		t.Fatalf("%#v", res)
	}
	sawToolErr := false
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.IsError {
			sawToolErr = true
		}
	}
	if !sawToolErr {
		t.Fatal("expected tool error to continue the loop")
	}
}

func TestLoopSteeringAndFollowUp(t *testing.T) {
	steered := false
	followed := false
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "list_dir", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "after-steer", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "after-follow", StopReason: llm.StopEnd},
		),
		Tools: tools.WorkspaceTools(t.TempDir()),
		GetSteering: func() []llm.Message {
			if steered {
				return nil
			}
			steered = true
			return []llm.Message{{Role: llm.RoleUser, Content: "add a footer"}}
		},
		GetFollowUp: func() []llm.Message {
			if followed {
				return nil
			}
			followed = true
			return []llm.Message{{Role: llm.RoleUser, Content: "one more"}}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Text != "after-follow" {
		t.Fatalf("text %q", res.Text)
	}
	joined := transcriptText(res.Messages)
	if !strings.Contains(joined, "add a footer") || !strings.Contains(joined, "one more") {
		t.Fatalf("transcript %s", joined)
	}
}

func TestLoopTranscriptReplay(t *testing.T) {
	path := filepath.Join(t.TempDir(), ".agentx", "transcript.jsonl")
	st, err := session.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	_, err = Run(context.Background(), Config{
		Model:        "fake",
		SystemPrompt: "be brief",
		Session:      st,
		Stream:       script(llm.Message{Role: llm.RoleAssistant, Content: "pong", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "ping"})
	if err != nil {
		t.Fatal(err)
	}
	loaded, err := st.Load()
	if err != nil {
		t.Fatal(err)
	}
	if len(loaded) != 3 || loaded[0].Role != llm.RoleSystem || loaded[1].Content != "ping" || loaded[2].Content != "pong" {
		t.Fatalf("%#v", loaded)
	}
	again, err := session.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	replay, err := again.Load()
	if err != nil {
		t.Fatal(err)
	}
	if len(replay) != len(loaded) || replay[2].Content != loaded[2].Content {
		t.Fatal("replay diverged")
	}
}

func TestLoopParallelReads(t *testing.T) {
	dir := t.TempDir()
	reg := tools.WorkspaceTools(dir)
	if reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"A"}`)).IsError {
		t.Fatal("seed a")
	}
	if reg.Execute(context.Background(), "write_file", []byte(`{"path":"b.txt","content":"B"}`)).IsError {
		t.Fatal("seed b")
	}
	res, err := Run(context.Background(), Config{
		Model:         "fake",
		Tools:         reg,
		ToolExecution: ExecParallel,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
				{ID: "1", Name: "read_file", Arguments: `{"path":"a.txt"}`},
				{ID: "2", Name: "read_file", Arguments: `{"path":"b.txt"}`},
			}},
			llm.Message{Role: llm.RoleAssistant, Content: "AB", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "read both"})
	if err != nil {
		t.Fatal(err)
	}
	got := map[string]string{}
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool {
			got[m.ToolCallID] = m.Content
		}
	}
	if got["1"] != "A" || got["2"] != "B" {
		t.Fatalf("parallel results %#v", got)
	}
}

func transcriptText(messages []llm.Message) string {
	var b strings.Builder
	for _, m := range messages {
		b.WriteString(m.Content)
		b.WriteByte('\n')
	}
	return b.String()
}