Agent-Native/runtime/r2_regression_test.go

package runtime

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

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

type countingTool struct {
	name  string
	calls int
	hook  func()
}

func (c *countingTool) Name() string        { return c.name }
func (c *countingTool) Description() string { return c.name }
func (c *countingTool) Parameters() json.RawMessage {
	return json.RawMessage(`{"type":"object","properties":{}}`)
}
func (c *countingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
	c.calls++
	if c.hook != nil {
		c.hook()
	}
	return tools.Result{Content: "ok"}, nil
}

func TestR2TokenBudgetSkipsTools(t *testing.T) {
	tool := &countingTool{name: "list_dir"}
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(tool)
	res, err := Run(context.Background(), Config{
		Model:  "fake",
		Tools:  reg,
		Budget: &Budget{MaxPromptTokens: 10},
		Stream: usageScript(llm.Usage{PromptTokens: 500},
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "list_dir", Arguments: `{}`,
			}}},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if tool.calls != 0 {
		t.Fatalf("tool ran after token budget: %d", tool.calls)
	}
	if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") {
		t.Fatalf("budget error: %v", res.Err)
	}
}

func TestR2StopReplacementIsPersisted(t *testing.T) {
	store, err := session.Open(filepath.Join(t.TempDir(), "t.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	_, err = Run(context.Background(), Config{
		Model:    "fake",
		MaxTurns: 3,
		Session:  store,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, Content: "draft", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "final", StopReason: llm.StopEnd},
		),
		Stop: func(last llm.Message, info StopInfo) StopDecision {
			if info.Continuations == 0 {
				return StopDecision{Continue: true, Replacement: "corrected-claim"}
			}
			return StopDecision{}
		},
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	msgs, err := store.Load()
	if err != nil {
		t.Fatal(err)
	}
	var saw bool
	for _, m := range msgs {
		if m.Content == "corrected-claim" {
			saw = true
		}
	}
	if !saw {
		t.Fatalf("replacement missing from session: %+v", msgs)
	}
}

func TestR2SequentialToolsStopAfterCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	first := &countingTool{name: "first", hook: cancel}
	second := &countingTool{name: "second"}
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(first)
	reg.Register(second)
	res, err := Run(ctx, Config{
		Model: "fake",
		Tools: reg,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
				{ID: "1", Name: "first", Arguments: `{}`},
				{ID: "2", Name: "second", Arguments: `{}`},
			}},
			llm.Message{Role: llm.RoleAssistant, Content: "after", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if first.calls != 1 || second.calls != 0 {
		t.Fatalf("calls first=%d second=%d", first.calls, second.calls)
	}
	var paired int
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool {
			paired++
		}
	}
	if paired != 2 {
		t.Fatalf("tool results = %d, want a pair for each call", paired)
	}
}