Cyrene_v2/internal/runtime/loop/mocks_test.go

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
}