Agent-Native/runtime/durable_test.go

package runtime

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

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

type replayTool struct {
	countingTool
	replay string
}

func (r *replayTool) Replay() string {
	if r.replay == "" {
		return tools.ReplayNever
	}
	return r.replay
}

func TestPendingUnsafeToolIsNotRerun(t *testing.T) {
	dir := t.TempDir()
	st, err := session.Open(filepath.Join(dir, "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	if err := st.Append(llm.Message{
		Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
		ToolCalls: []llm.ToolCall{{ID: "c1", Name: "write_file", Arguments: `{}`}},
	}); err != nil {
		t.Fatal(err)
	}
	if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "write_file", Arguments: `{}`, Replay: tools.ReplayNever}); err != nil {
		t.Fatal(err)
	}
	tool := &countingTool{name: "write_file"}
	reg := tools.NewRegistry(dir)
	reg.Register(tool)
	_, err = Run(context.Background(), Config{
		Model: "fake", Session: st, Tools: reg, MaxTurns: 2,
		Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "next"})
	if err != nil {
		t.Fatal(err)
	}
	if tool.calls != 0 {
		t.Fatalf("unsafe tool reran: %d", tool.calls)
	}
	msgs, err := st.Load()
	if err != nil {
		t.Fatal(err)
	}
	var paired bool
	for _, m := range msgs {
		if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.Content == unknownSideEffect {
			paired = true
		}
	}
	if !paired {
		t.Fatalf("pair missing: %#v", msgs)
	}
}

func TestPendingSafeToolRerunsOnce(t *testing.T) {
	dir := t.TempDir()
	st, err := session.Open(filepath.Join(dir, "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	if err := st.Append(llm.Message{
		Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
		ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}},
	}); err != nil {
		t.Fatal(err)
	}
	if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil {
		t.Fatal(err)
	}
	tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplaySafe}
	reg := tools.NewRegistry(dir)
	reg.Register(tool)
	_, err = Run(context.Background(), Config{
		Model: "fake", Session: st, Tools: reg, MaxTurns: 2,
		Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "next"})
	if err != nil {
		t.Fatal(err)
	}
	if tool.calls != 1 {
		t.Fatalf("calls %d", tool.calls)
	}
	msgs, err := st.Load()
	if err != nil {
		t.Fatal(err)
	}
	for _, m := range msgs {
		if m.Role == llm.RoleTool && m.ToolCallID == "c1" {
			if m.Content != "ok" || strings.Contains(m.Content, unknownSideEffect) {
				t.Fatalf("result %#v", m)
			}
			return
		}
	}
	t.Fatal("missing tool result")
}

func TestStoredSafeCurrentNeverDoesNotRerun(t *testing.T) {
	dir := t.TempDir()
	st, err := session.Open(filepath.Join(dir, "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil {
		t.Fatal(err)
	}
	tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplayNever}
	reg := tools.NewRegistry(dir)
	reg.Register(tool)
	if err := recoverPending(context.Background(), Config{Session: st, Tools: reg}); err != nil {
		t.Fatal(err)
	}
	if tool.calls != 0 {
		t.Fatalf("calls %d", tool.calls)
	}
}

func TestFollowUpIsPersistedBeforeNextRequest(t *testing.T) {
	dir := t.TempDir()
	st, err := session.Open(filepath.Join(dir, "s.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	sent := false
	var saw bool
	stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 1)
		go func() {
			defer close(ch)
			for _, m := range messages {
				if m.Content == "follow-please" {
					loaded, err := st.Load()
					if err != nil {
						t.Errorf("load %v", err)
					}
					for _, got := range loaded {
						if got.Content == "follow-please" {
							saw = true
						}
					}
				}
			}
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
		}()
		return ch
	}
	_, err = Run(context.Background(), Config{
		Model: "fake", Session: st, MaxTurns: 3, Stream: stream,
		GetFollowUp: func() []llm.Message {
			if sent {
				return nil
			}
			sent = true
			return []llm.Message{{Role: llm.RoleUser, Content: "follow-please"}}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if !saw {
		t.Fatal("follow-up entered the request before it was stored")
	}
}

func TestSpansUseAgreedNamesWithoutPrompt(t *testing.T) {
	var names []string
	var leaked bool
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(&countingTool{name: "read_file"})
	_, err := Run(context.Background(), Config{
		Model: "fake", Tools: reg, MaxTurns: 2,
		Trace: func(name string, attrs map[string]string) func() {
			names = append(names, name)
			if _, ok := attrs["prompt"]; ok || attrs["content"] != "" {
				leaked = true
			}
			return func() {}
		},
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "1", Name: "read_file", Arguments: `{}`}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "用户的秘密提示词"})
	if err != nil {
		t.Fatal(err)
	}
	joined := strings.Join(names, ",")
	for _, want := range []string{"invoke_agent", "chat", "execute_tool"} {
		if !strings.Contains(joined, want) {
			t.Fatalf("names %s", joined)
		}
	}
	if leaked {
		t.Fatal("prompt leaked into span attrs")
	}
}