Agent-Native/runtime/ab_gate_test.go

package runtime

import (
	"context"
	"errors"
	"testing"

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

func TestHookAdaptersMatchDirectCallbacks(t *testing.T) {
	direct := runGate(t, false)
	adapted := runGate(t, true)
	if direct.stop != adapted.stop || direct.text != adapted.text {
		t.Fatalf("stop/text direct=%s %q adapted=%s %q", direct.stop, direct.text, adapted.stop, adapted.text)
	}
	if direct.denies != adapted.denies || direct.reminders != adapted.reminders || direct.errors != adapted.errors {
		t.Fatalf("counts direct=%+v adapted=%+v", direct, adapted)
	}
	if direct.denies != 1 || direct.reminders == 0 {
		t.Fatalf("unexpected counts %+v", direct)
	}
}

type gateStats struct {
	stop, text                string
	denies, reminders, errors int
}

func runGate(t *testing.T, viaHooks bool) gateStats {
	t.Helper()
	var stats gateStats
	before := func(call llm.ToolCall) error {
		if call.Name == "read_file" {
			stats.denies++
			return errors.New("denied")
		}
		return nil
	}
	reminders := func(int) []string {
		stats.reminders++
		return []string{"记一笔"}
	}
	stop := func(last llm.Message, info StopInfo) StopDecision {
		if info.Continuations < 1 {
			return StopDecision{Continue: true, Reason: "再看一眼"}
		}
		return StopDecision{}
	}
	cfg := Config{
		Model: "fake",
		Tools: tools.WorkspaceTools(t.TempDir()),
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "read_file", Arguments: `{"path":"a.md"}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
		Reminders: reminders,
	}
	if viaHooks {
		bus := hooks.New()
		bus.On(hooks.AdaptBefore(func(name, args string) error {
			return before(llm.ToolCall{Name: name, Arguments: args})
		}))
		bus.On(func(c hooks.Call) hooks.Outcome {
			if c.Event != hooks.Stop {
				return hooks.Outcome{}
			}
			d := stop(llm.Message{Content: c.LastText}, StopInfo{Continuations: c.Continuations, ToolCalls: c.ToolCalls})
			if d.Continue {
				return hooks.Outcome{Decision: hooks.Block, Reason: d.Reason, Replacement: d.Replacement, Notice: d.Notice}
			}
			return hooks.Outcome{Decision: hooks.Allow, Notice: d.Notice}
		})
		cfg.Hooks = bus
	} else {
		cfg.BeforeTool = before
		cfg.Stop = stop
	}
	res, err := Run(context.Background(), cfg, llm.Message{Role: llm.RoleUser, Content: "读文件"})
	if err != nil {
		t.Fatal(err)
	}
	stats.stop = string(res.StopReason)
	stats.text = res.Text
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.IsError {
			stats.errors++
		}
	}
	return stats
}