Cyrene_v2/internal/runtime/loop/loop_extra_test.go

package loop

import (
	"context"
	"strings"
	"sync"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

func TestRun_DuplicateCallInBatch_Skipped(t *testing.T) {
	t.Parallel()
	// Two identical calls in one batch: the second must be skipped, not executed,
	// and fed back as a duplicate hint.
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}},{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "done" {
		t.Fatalf("Text = %q, want done", res.Text)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1 (duplicate skipped)", res.ToolCallsExecuted)
	}
	if res.StopReason != agent.StopReasonStop {
		t.Fatalf("StopReason = %q, want stop", res.StopReason)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	// Round 1 appended two tool messages: the executed call result, then the
	// duplicate skip. The last two messages before the final turn are those.
	result := last[len(last)-2]
	skip := last[len(last)-1]
	if !strings.Contains(result.Content[0].Text, "hi") {
		t.Fatalf("expected echo result, got %q", result.Content[0].Text)
	}
	if !strings.Contains(skip.Content[0].Text, "skipped") || !strings.Contains(skip.Content[0].Text, "duplicate") {
		t.Fatalf("expected duplicate-skip message, got %q", skip.Content[0].Text)
	}
}

func TestRun_CriticalToolLoop_Terminates(t *testing.T) {
	t.Parallel()
	// The identical call repeats across rounds: round 2 is blocked (critical
	// loop flagged), round 3 terminates the run.
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: call}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.StopReason != StopReasonToolLoopTerminated {
		t.Fatalf("StopReason = %q, want %q", res.StopReason, StopReasonToolLoopTerminated)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1 (only round 1 executed)", res.ToolCallsExecuted)
	}
	if res.Steps != 3 {
		t.Fatalf("Steps = %d, want 3", res.Steps)
	}
	if !strings.Contains(res.Text, "critical tool-loop") {
		t.Fatalf("Text = %q, want critical tool-loop message", res.Text)
	}
}

func TestRun_Parallel_ResultsInCallOrder(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"aa","args":{}},{"name":"bb","args":{}},{"name":"cc","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "final"}},
	}}
	td := dispatcherWith(t,
		&mockTool{name: "aa", result: "va"},
		&mockTool{name: "bb", result: "vb"},
		&mockTool{name: "cc", result: "vc"},
	)
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, ToolExecution: "parallel"})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "final" {
		t.Fatalf("Text = %q, want final", res.Text)
	}
	if res.ToolCallsExecuted != 3 {
		t.Fatalf("ToolCallsExecuted = %d, want 3", res.ToolCallsExecuted)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	if len(last) < 4 {
		t.Fatalf("messages = %d, want at least 4 (user + 3 results)", len(last))
	}
	got := make([]string, 3)
	for i := range got {
		got[i] = last[len(last)-3+i].Content[0].Text
	}
	want := []string{
		"tool aa result: va",
		"tool bb result: vb",
		"tool cc result: vc",
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("result %d = %q, want %q (call order must be preserved)", i, got[i], want[i])
		}
	}
}

func TestRun_Sequential_ExecutesAll(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"aa","args":{}},{"name":"bb","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "final"}},
	}}
	td := dispatcherWith(t,
		&mockTool{name: "aa", result: "va"},
		&mockTool{name: "bb", result: "vb"},
	)
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, ToolExecution: "sequential"})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 2 {
		t.Fatalf("ToolCallsExecuted = %d, want 2", res.ToolCallsExecuted)
	}
}

func TestRun_ArgumentValidation_ErrorKind(t *testing.T) {
	t.Parallel()
	// Missing required field -> the fed-back failure carries errorKind
	// argument-validation so the model can re-issue a corrected call.
	call := `{"tool_calls":[{"name":"req","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "recovered"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "req", params: map[string]any{"required": []any{"must"}}})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "recovered" {
		t.Fatalf("Text = %q, want recovered", res.Text)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	lastMsg := last[len(last)-1]
	if !strings.Contains(lastMsg.Content[0].Text, `"kind":"argument-validation"`) {
		t.Fatalf("result should carry argument-validation kind, got %q", lastMsg.Content[0].Text)
	}
}

func TestRun_BatchAdmit_BlocksBatch(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "after block"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, Hooks: ToolLoopHooks{
		BatchAdmit: func(ctx context.Context, calls []agent.ToolCall) (bool, string) {
			return false, "policy denied"
		},
	}})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0 (batch blocked)", res.ToolCallsExecuted)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	lastMsg := last[len(last)-1]
	if !strings.Contains(lastMsg.Content[0].Text, "policy denied") {
		t.Fatalf("expected blocked batch message, got %q", lastMsg.Content[0].Text)
	}
}

func TestRun_CallAdmit_VetoesSingleCall(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "after veto"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, Hooks: ToolLoopHooks{
		CallAdmit: func(ctx context.Context, call agent.ToolCall) *ToolCallVeto {
			if call.Name == "echo" {
				return &ToolCallVeto{Block: true, Reason: "sensitive action requires approval"}
			}
			return nil
		},
	}})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0 (call vetoed)", res.ToolCallsExecuted)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	lastMsg := last[len(last)-1]
	if !strings.Contains(lastMsg.Content[0].Text, "sensitive action requires approval") {
		t.Fatalf("expected veto message, got %q", lastMsg.Content[0].Text)
	}
}

func TestRun_AfterCall_RewritesResult(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "final"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "raw"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, Hooks: ToolLoopHooks{
		AfterCall: func(ctx context.Context, call agent.ToolCall, outcome ToolOutcome) *ToolOutcome {
			if call.Name == "echo" {
				outcome.Result = "rewritten"
				return &outcome
			}
			return nil
		},
	}})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	lastMsg := last[len(last)-1]
	if !strings.Contains(lastMsg.Content[0].Text, "rewritten") {
		t.Fatalf("expected rewritten result, got %q", lastMsg.Content[0].Text)
	}
}

func TestRun_ToolCallRepair_FencedJSON(t *testing.T) {
	t.Parallel()
	// A code-fenced tool-call envelope is repaired and executed rather than
	// silently becoming a text answer.
	call := "```json\n{\"tool_calls\":[{\"name\":\"echo\",\"args\":{\"q\":\"hi\"}}]}\n```"
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 1 || res.Text != "done" {
		t.Fatalf("ToolCallsExecuted = %d, Text = %q; want 1 + done (fenced JSON repaired)", res.ToolCallsExecuted, res.Text)
	}
}

func TestRun_StreamToolExecution_EmitsProgress(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{
		stream: []mgResp{
			{resp: model.Response{Content: call}},
			{resp: model.Response{Content: "done"}},
		},
		chunks: [][]string{{"x"}},
	}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	var mu sync.Mutex
	var evs []model.StreamEvent
	l := New(Options{
		Agent: a,
		OnStream: func(ev model.StreamEvent) {
			mu.Lock()
			evs = append(evs, ev)
			mu.Unlock()
		},
		MaxRounds: 8,
	})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	mu.Lock()
	defer mu.Unlock()
	saw := false
	for _, ev := range evs {
		if ev.Kind == model.StreamKindToolExecution {
			saw = true
			break
		}
	}
	if !saw {
		t.Fatalf("expected a tool_execution_update event on the streaming path, got %+v", evs)
	}
}