feat(loop): agent loop（8 轮 + checkpoint + 流式 + 记忆注入/提取 + trace + 工具超时回喂）

2026-09-06 05:33:38 +08:00
1 parent 01a4e29f84
commit 71697ca4f8
@@ -0,0 +1,204 @@
+package loop
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"testing"
+	"time"
+
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
+	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
+)
+
+// mgResp is one scripted adapter result: a response plus an optional error.
+type mgResp struct {
+	resp model.Response
+	err  error
+}
+
+// mockAdapter is a configurable model.Adapter for loop tests. It keeps two
+// separate scripts and call counters (one for Chat, one for ChatStream) so a
+// single instance serves both the non-streaming and streaming paths. Every
+// request it is handed is recorded so tests can assert on the messages that were
+// fed to the model (e.g. the injected memory or a fed-back tool result).
+type mockAdapter struct {
+	name string
+
+	chat      []mgResp
+	chatCalls int
+	reqs      []model.Request
+
+	stream      []mgResp
+	chunks      [][]string
+	streamCalls int
+	streamReqs  []model.Request
+}
+
+func (m *mockAdapter) Name() string {
+	if m.name == "" {
+		return "mock"
+	}
+	return m.name
+}
+
+func (m *mockAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
+	m.reqs = append(m.reqs, req)
+	if m.chatCalls >= len(m.chat) {
+		return model.Response{Content: "done"}, nil
+	}
+	r := m.chat[m.chatCalls]
+	m.chatCalls++
+	return r.resp, r.err
+}
+
+func (m *mockAdapter) ChatStream(_ context.Context, req model.Request, on func(string)) (model.Response, error) {
+	m.streamReqs = append(m.streamReqs, req)
+	i := m.streamCalls
+	m.streamCalls++
+	if i < len(m.chunks) {
+		for _, c := range m.chunks[i] {
+			if on != nil {
+				on(c)
+			}
+		}
+	}
+	if i >= len(m.stream) {
+		return model.Response{Content: "done"}, nil
+	}
+	return m.stream[i].resp, m.stream[i].err
+}
+
+// mockGuard is a configurable guardrail.Guardrail for loop tests.
+type mockGuard struct {
+	inputVerdict  *guardrail.Verdict
+	outputVerdict *guardrail.Verdict
+	inputErr      error
+	outputErr     error
+}
+
+func (m *mockGuard) CheckInput(_ context.Context, _ *model.Request) (*guardrail.Verdict, error) {
+	return m.inputVerdict, m.inputErr
+}
+
+func (m *mockGuard) CheckOutput(_ context.Context, _ *model.Response) (*guardrail.Verdict, error) {
+	return m.outputVerdict, m.outputErr
+}
+
+// mockTool is a minimal sdk.Tool for loop tests. It supports a delay so a test
+// can exercise the per-invoke timeout, and an error so a test can verify that a
+// failed call is fed back to the model.
+type mockTool struct {
+	name   string
+	desc   string
+	params map[string]any
+	result any
+	err    error
+	delay  time.Duration
+}
+
+func (m *mockTool) Name() string { return m.name }
+
+func (m *mockTool) Description() string {
+	if m.desc == "" {
+		return "mock tool " + m.name
+	}
+	return m.desc
+}
+
+func (m *mockTool) Parameters() (map[string]any, error) { return m.params, nil }
+
+func (m *mockTool) Execute(ctx context.Context, _ any) (any, error) {
+	if m.delay > 0 {
+		select {
+		case <-time.After(m.delay):
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+	return m.result, m.err
+}
+
+// testLogger returns a logger that writes nowhere, keeping test output clean.
+func testLogger() *slog.Logger {
+	return slog.New(slog.NewTextHandler(io.Discard, nil))
+}
+
+// newTestAgent builds an Agent over ad with the given tools and no guardrails.
+func newTestAgent(t *testing.T, ad model.Adapter, tools *tool.Dispatcher) *agent.Agent {
+	t.Helper()
+	a, err := agent.New("You are a helpful assistant.", tools, nil, agent.WithModel(ad, nil))
+	if err != nil {
+		t.Fatalf("agent.New: %v", err)
+	}
+	return a
+}
+
+// dispatcherWith builds a Dispatcher with the default tool timeout.
+func dispatcherWith(t *testing.T, tools ...sdk.Tool) *tool.Dispatcher {
+	t.Helper()
+	return dispatcherWithTimeout(t, 0, tools...)
+}
+
+// dispatcherWithTimeout builds a Dispatcher applying a custom default timeout
+// (a non-positive value selects the package default) and registering tools.
+func dispatcherWithTimeout(t *testing.T, timeout time.Duration, tools ...sdk.Tool) *tool.Dispatcher {
+	t.Helper()
+	reg := registry.New()
+	var opts []tool.Option
+	if timeout > 0 {
+		opts = append(opts, tool.WithDefaultTimeout(timeout))
+	}
+	d, err := tool.New(reg, opts...)
+	if err != nil {
+		t.Fatalf("tool.New: %v", err)
+	}
+	for _, tk := range tools {
+		if err := d.Register(tk); err != nil {
+			t.Fatalf("Register %q: %v", tk.Name(), err)
+		}
+	}
+	return d
+}
+
+// newTestMemoryStore opens an isolated in-memory SQLite database and returns a
+// memory.Store over it. Each call owns a private DB, so parallel tests are
+// isolated by construction.
+func newTestMemoryStore(t *testing.T) *memory.Store {
+	t.Helper()
+	ctx := context.Background()
+	d, err := db.Open(ctx, ":memory:", testLogger())
+	if err != nil {
+		t.Fatalf("db.Open: %v", err)
+	}
+	t.Cleanup(func() { _ = d.Close() })
+	s, err := memory.NewStore(ctx, d, testLogger())
+	if err != nil {
+		t.Fatalf("memory.NewStore: %v", err)
+	}
+	return s
+}
+
+// newTestCheckpointStore opens an isolated in-memory SQLite database and returns
+// a persist.Store over it.
+func newTestCheckpointStore(t *testing.T) *persist.Store {
+	t.Helper()
+	ctx := context.Background()
+	d, err := db.Open(ctx, ":memory:", testLogger())
+	if err != nil {
+		t.Fatalf("db.Open: %v", err)
+	}
+	t.Cleanup(func() { _ = d.Close() })
+	s, err := persist.New(ctx, d, testLogger())
+	if err != nil {
+		t.Fatalf("persist.New: %v", err)
+	}
+	return s
+}