byteclaw/internal/agent/agent_test.go

package agent

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/crypto/byteclaw/internal/llm"
	"github.com/crypto/byteclaw/internal/store"
	"github.com/crypto/byteclaw/internal/tools"
)

// scripted provider returns queued responses and records requests.
type scripted struct {
	resps []llm.Response
	reqs  []llm.Request
}

func (s *scripted) Chat(_ context.Context, req llm.Request, d llm.DeltaFunc) (llm.Response, error) {
	s.reqs = append(s.reqs, req)
	r := s.resps[0]
	s.resps = s.resps[1:]
	if d != nil && r.Content != "" {
		d(r.Content)
	}
	return r, nil
}

func newStore(t *testing.T) *store.Store {
	s, err := store.Open(filepath.Join(t.TempDir(), "a.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	return s
}

func TestToolLoop(t *testing.T) {
	st := newStore(t)
	ss, _ := st.CreateSession(context.Background(), "")
	p := &scripted{resps: []llm.Response{
		{ToolCalls: []llm.ToolCall{{ID: "c1", Name: "time_now", Arguments: `{"timezone":"UTC"}`}}},
		{Content: "It is noon."},
	}}
	fixed := time.Date(2026, 5, 5, 12, 0, 0, 0, time.UTC)
	a := &Agent{LLM: p, Store: st, Tools: tools.NewRegistry(tools.TimeNow{Now: func() time.Time { return fixed }}), SystemPrompt: "sys", MaxIterations: 4}
	var events []Event
	out, err := a.Run(context.Background(), ss.ID, "what time?", func(e Event) { events = append(events, e) })
	if err != nil || out != "It is noon." {
		t.Fatalf("%q %v", out, err)
	}
	second := p.reqs[1].Messages
	if second[0].Role != llm.RoleSystem || second[len(second)-1].Role != llm.RoleTool || !strings.HasPrefix(second[len(second)-1].Content, "2026-05-05T12:00:00") {
		t.Fatalf("bad second request: %+v", second)
	}
	types := []string{}
	for _, e := range events {
		types = append(types, e.Type)
	}
	if strings.Join(types, ",") != "tool_call,tool_result,delta" {
		t.Fatalf("events %v", types)
	}
	msgs, _ := st.Messages(context.Background(), ss.ID, 0)
	if len(msgs) != 4 {
		t.Fatalf("stored %d msgs", len(msgs))
	}
}

func TestMaxIterations(t *testing.T) {
	st := newStore(t)
	ss, _ := st.CreateSession(context.Background(), "")
	call := llm.Response{ToolCalls: []llm.ToolCall{{ID: "c", Name: "time_now", Arguments: "{}"}}}
	a := &Agent{LLM: &scripted{resps: []llm.Response{call, call}}, Store: st, Tools: tools.NewRegistry(tools.TimeNow{}), MaxIterations: 2}
	if _, err := a.Run(context.Background(), ss.ID, "loop", nil); err == nil {
		t.Fatal("expected iteration error")
	}
}

func TestHistoryDropsOrphanToolMessages(t *testing.T) {
	st := newStore(t)
	ctx := context.Background()
	ss, _ := st.CreateSession(ctx, "")
	st.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "x", Name: "t"}}})
	st.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleTool, ToolCallID: "x", Content: "r"})
	st.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleAssistant, Content: "ok"})
	a := &Agent{Store: st, HistoryLimit: 2}
	h, _ := a.history(ctx, ss.ID)
	if len(h) != 1 || h[0].Content != "ok" {
		t.Fatalf("%+v", h)
	}
}

// blocking streams one delta then waits for cancellation, like a slow provider.
type blocking struct{ started chan struct{} }

func (b *blocking) Chat(ctx context.Context, _ llm.Request, d llm.DeltaFunc) (llm.Response, error) {
	d("Teil")
	d("antwort")
	close(b.started)
	<-ctx.Done()
	return llm.Response{}, ctx.Err()
}

func TestCancelStoresPartialAnswer(t *testing.T) {
	st := newStore(t)
	ss, _ := st.CreateSession(context.Background(), "")
	b := &blocking{started: make(chan struct{})}
	a := &Agent{LLM: b, Store: st, Tools: tools.NewRegistry(), MaxIterations: 3}
	if a.Cancel(ss.ID) {
		t.Fatal("nothing should be running")
	}
	go func() {
		<-b.started
		if !a.Cancel(ss.ID) {
			t.Error("cancel should find the run")
		}
	}()
	text, err := a.Run(context.Background(), ss.ID, "erzähl was", nil)
	if !errors.Is(err, ErrCancelled) || text != "Teilantwort" {
		t.Fatalf("%q %v", text, err)
	}
	msgs, _ := st.Messages(context.Background(), ss.ID, 0)
	if len(msgs) != 2 || !msgs[1].Cancelled || msgs[1].Content != "Teilantwort" {
		t.Fatalf("%+v", msgs)
	}
	if a.Running(ss.ID) {
		t.Fatal("run must be unregistered")
	}
}

// recorder answers summary requests with "SUMMARY" and normal turns with "ok".
type recorder struct{ reqs []llm.Request }

func (r *recorder) Chat(_ context.Context, req llm.Request, _ llm.DeltaFunc) (llm.Response, error) {
	r.reqs = append(r.reqs, req)
	if strings.Contains(req.Messages[0].Content, "compress chat histories") {
		return llm.Response{Content: "SUMMARY"}, nil
	}
	return llm.Response{Content: "ok"}, nil
}

func TestAutoSummarize(t *testing.T) {
	st := newStore(t)
	ctx := context.Background()
	ss, _ := st.CreateSession(ctx, "")
	p := &recorder{}
	a := &Agent{LLM: p, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2, SystemPrompt: "sys",
		Summarize: SummarizeConfig{Enabled: true, MaxMessages: 6, KeepRecent: 2}}
	for i := 0; i < 3; i++ { // 6 messages -> at threshold, no summary yet
		a.Run(ctx, ss.ID, fmt.Sprintf("frage %d", i), nil)
	}
	if _, err := st.LatestSummary(ctx, ss.ID); err == nil {
		t.Fatal("summary created too early")
	}
	var events []string
	a.Run(ctx, ss.ID, "frage 3", func(e Event) { events = append(events, e.Type) })
	sm, err := st.LatestSummary(ctx, ss.ID)
	if err != nil || sm.Content != "SUMMARY" {
		t.Fatalf("%+v %v", sm, err)
	}
	if events[0] != "summarized" {
		t.Fatalf("%v", events)
	}
	last := p.reqs[len(p.reqs)-1].Messages
	if last[1].Role != llm.RoleSystem || !strings.Contains(last[1].Content, "SUMMARY") {
		t.Fatalf("summary not in context: %+v", last)
	}
	// the verbatim tail starts at a user message and older messages are not sent
	if last[2].Role != llm.RoleUser || strings.Contains(fmt.Sprint(last), "frage 0") {
		t.Fatalf("bad window: %+v", last)
	}
	if all, _ := st.Messages(ctx, ss.ID, 0); len(all) != 8 {
		t.Fatalf("originals must be kept, got %d", len(all))
	}
}

func TestAutoSummarizeByTokens(t *testing.T) {
	st := newStore(t)
	ctx := context.Background()
	ss, _ := st.CreateSession(ctx, "")
	a := &Agent{LLM: &recorder{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2,
		Summarize: SummarizeConfig{Enabled: true, MaxMessages: 1000, MaxTokens: 50, KeepRecent: 1}}
	a.Run(ctx, ss.ID, strings.Repeat("lang ", 100), nil)
	a.Run(ctx, ss.ID, "kurz", nil)
	if _, err := st.LatestSummary(ctx, ss.ID); err != nil {
		t.Fatal("token threshold should trigger summary")
	}
}

func TestSummarizeNowManual(t *testing.T) {
	st := newStore(t)
	ctx := context.Background()
	ss, _ := st.CreateSession(ctx, "")
	a := &Agent{LLM: &recorder{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2, Summarize: SummarizeConfig{KeepRecent: 2}}
	if sm, err := a.SummarizeNow(ctx, ss.ID); err != nil || sm != nil {
		t.Fatalf("empty session: %v %v", sm, err)
	}
	a.Run(ctx, ss.ID, "eins", nil)
	a.Run(ctx, ss.ID, "zwei", nil)
	sm, err := a.SummarizeNow(ctx, ss.ID)
	if err != nil || sm == nil || sm.Content != "SUMMARY" {
		t.Fatalf("%v %v", sm, err)
	}
}